Let the Environment override the CodeGenerator

see #404

diff --git a/docs/api.rst b/docs/api.rst
index d143ebc59e40a16b9baa1a0ed38b47b17212d95f..aaf262b34d6a371b6cca1d78969b73d1e69bc2d4 100644 (file)
--- 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 96b32f7d6b4af673371c1a8dbd5298a6151cd6f3..5cd872dfdf31e0e21bad3ba45fe528c5f90b7d5c 100644 (file)
--- 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 1f6496b786c421e5c2322c9d08ab7c2f39d9d849..3bd8fa7b49a3cbdc38f288aa43734d1a750c6b95 100644 (file)
--- 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 40a6b3dc9d1d67b470a1859ed14341b9e5fdc5c1..5cd7f86f25d1afdccb8072f7c6f8eff2727bd120 100644 (file)
--- 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'