From 0c1f192670fff219ae1c61fb69c75ae4fb0fd50d Mon Sep 17 00:00:00 2001
From: Armin Ronacher <armin.ronacher@active-4.com>
Date: Mon, 25 May 2015 12:24:09 +0200
Subject: [PATCH] Added decorator support for finalize.  This fixes #431

---
 CHANGES            |  2 ++
 jinja2/compiler.py | 26 ++++++++++++++++++++++++--
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/CHANGES b/CHANGES
index 241a95c8..cfe4c43c 100644
--- a/CHANGES
+++ b/CHANGES
@@ -29,6 +29,8 @@ Version 2.8
 - Add ability to use custom subclasses of ``jinja2.compiler.CodeGenerator`` and
   ``jinja2.runtime.Context`` by adding two new attributes to the environment
   (`code_generator_class` and `context_class`) (pull request ``#404``).
+- added support for context/environment/evalctx decorator functions on
+  the finalize callback of the environment.
 
 Version 2.7.3
 -------------
diff --git a/jinja2/compiler.py b/jinja2/compiler.py
index 5cd872df..fad007b5 100644
--- a/jinja2/compiler.py
+++ b/jinja2/compiler.py
@@ -1219,8 +1219,17 @@ class CodeGenerator(NodeVisitor):
         if self.has_known_extends and frame.require_output_check:
             return
 
+        allow_constant_finalize = True
         if self.environment.finalize:
-            finalize = lambda x: text_type(self.environment.finalize(x))
+            func = self.environment.finalize
+            if getattr(func, 'contextfunction', False) or \
+               getattr(func, 'evalcontextfunction', False):
+                allow_constant_finalize = False
+            elif getattr(func, 'environmentfunction', False):
+                finalize = lambda x: text_type(
+                    self.environment.finalize(self.environment, x))
+            else:
+                finalize = lambda x: text_type(self.environment.finalize(x))
         else:
             finalize = text_type
 
@@ -1236,6 +1245,8 @@ class CodeGenerator(NodeVisitor):
         body = []
         for child in node.nodes:
             try:
+                if not allow_constant_finalize:
+                    raise nodes.Impossible()
                 const = child.as_const(frame.eval_ctx)
             except nodes.Impossible:
                 body.append(child)
@@ -1291,6 +1302,9 @@ class CodeGenerator(NodeVisitor):
                         self.write('to_string(')
                     if self.environment.finalize is not None:
                         self.write('environment.finalize(')
+                        if getattr(self.environment.finalize,
+                                   "contextfunction", False):
+                            self.write('context, ')
                         close += 1
                     self.visit(item, frame)
                     self.write(')' * close)
@@ -1313,7 +1327,6 @@ class CodeGenerator(NodeVisitor):
                     arguments.append(item)
             self.writeline('yield ')
             self.write(repr(concat(format)) + ' % (')
-            idx = -1
             self.indent()
             for argument in arguments:
                 self.newline(argument)
@@ -1327,6 +1340,15 @@ class CodeGenerator(NodeVisitor):
                     close += 1
                 if self.environment.finalize is not None:
                     self.write('environment.finalize(')
+                    if getattr(self.environment.finalize,
+                               'contextfunction', False):
+                        self.write('context, ')
+                    elif getattr(self.environment.finalize,
+                               'evalcontextfunction', False):
+                        self.write('context.eval_ctx, ')
+                    elif getattr(self.environment.finalize,
+                               'environmentfunction', False):
+                        self.write('environment, ')
                     close += 1
                 self.visit(argument, frame)
                 self.write(')' * close + ', ')
-- 
2.47.2