From: ThiefMaster <adrian@planetcoding.net>
Date: Mon, 6 Apr 2015 11:54:14 +0000 (+0200)
Subject: Let the Environment override the CodeGenerator
X-Git-Tag: 2.8~18^2~2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=14936312c0755f811e73e691ab495e39afc6b7d8;p=thirdparty%2Fjinja.git

Let the Environment override the CodeGenerator

see #404
---

diff --git a/docs/api.rst b/docs/api.rst
index d143ebc5..aaf262b3 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -154,6 +154,12 @@ useful if you want to dig deeper into Jinja2 or :ref:`develop extensions
         to modify this dict.  For more details see :ref:`global-namespace`.
         For valid object names have a look at :ref:`identifier-naming`.
 
+    .. attribute:: code_generator_class
+
+       The class used for code generation.  This should not be changed
+       in most cases, unless you need to modify the Python code a
+       template compiles to.
+
     .. automethod:: overlay([options])
 
     .. method:: undefined([hint, obj, name, exc])
diff --git a/jinja2/compiler.py b/jinja2/compiler.py
index 96b32f7d..5cd872df 100644
--- a/jinja2/compiler.py
+++ b/jinja2/compiler.py
@@ -57,7 +57,8 @@ def generate(node, environment, name, filename, stream=None,
     """Generate the python source for a node tree."""
     if not isinstance(node, nodes.Template):
         raise TypeError('Can\'t compile non template nodes')
-    generator = CodeGenerator(environment, name, filename, stream, defer_init)
+    generator = environment.code_generator_class(environment, name, filename,
+                                                 stream, defer_init)
     generator.visit(node)
     if stream is None:
         return generator.stream.getvalue()
diff --git a/jinja2/environment.py b/jinja2/environment.py
index 1f6496b7..3bd8fa7b 100644
--- a/jinja2/environment.py
+++ b/jinja2/environment.py
@@ -21,7 +21,7 @@ from jinja2.lexer import get_lexer, TokenStream
 from jinja2.parser import Parser
 from jinja2.nodes import EvalContext
 from jinja2.optimizer import optimize
-from jinja2.compiler import generate
+from jinja2.compiler import generate, CodeGenerator
 from jinja2.runtime import Undefined, new_context
 from jinja2.exceptions import TemplateSyntaxError, TemplateNotFound, \
      TemplatesNotFound, TemplateRuntimeError
@@ -238,6 +238,10 @@ class Environment(object):
     exception_handler = None
     exception_formatter = None
 
+    #: the class that is used for code generation.  See
+    #: :class:`~jinja2.compiler.CodeGenerator` for more information.
+    code_generator_class = CodeGenerator
+
     def __init__(self,
                  block_start_string=BLOCK_START_STRING,
                  block_end_string=BLOCK_END_STRING,
diff --git a/tests/test_api.py b/tests/test_api.py
index 40a6b3dc..5cd7f86f 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -16,6 +16,7 @@ import pytest
 from jinja2 import Environment, Undefined, DebugUndefined, \
      StrictUndefined, UndefinedError, meta, \
      is_undefined, Template, DictLoader, make_logging_undefined
+from jinja2.compiler import CodeGenerator
 from jinja2.utils import Cycler
 
 
@@ -290,3 +291,24 @@ class TestUndefined():
             assert e.message == "'int object' has no attribute 'upper'"
         else:
             assert False, 'expected exception'
+
+
+@pytest.mark.api
+@pytest.mark.lowlevel
+class TestLowLevel():
+
+    def test_custom_code_generator(self):
+        class CustomCodeGenerator(CodeGenerator):
+            def visit_Const(self, node, frame=None):
+                # This method is pure nonsense, but works fine for testing...
+                if node.value == 'foo':
+                    self.write(repr('bar'))
+                else:
+                    super(CustomCodeGenerator, self).visit_Const(node, frame)
+
+        class CustomEnvironment(Environment):
+            code_generator_class = CustomCodeGenerator
+
+        env = CustomEnvironment()
+        tmpl = env.from_string('{% set foo = "foo" %}{{ foo }}')
+        assert tmpl.render() == 'bar'