Merge branch 'override-codegen-context'

CHANGES

@@ -26,6 +26,9 @@ Version 2.8
   change makes ``{% macro m(x, y=1, z) %}...{% endmacro %}`` a syntax error. The
   previous behavior for this code was broken anyway (resulting in the default
   value being applied to `y`).
+- 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``).
 
 Version 2.7.3
 -------------

docs/api.rst

@@ -154,6 +154,19 @@ 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.
+
+    .. attribute:: context_class
+
+       The context used for templates.  This should not be changed
+       in most cases, unless you need to modify internals of how
+       template variables are handled.  For details, see
+       :class:`~jinja2.runtime.Context`.
+
     .. automethod:: overlay([options])
 
     .. method:: undefined([hint, obj, name, exc])

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()

jinja2/environment.py

@@ -21,8 +21,8 @@ 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.runtime import Undefined, new_context
+from jinja2.compiler import generate, CodeGenerator
+from jinja2.runtime import Undefined, new_context, Context
 from jinja2.exceptions import TemplateSyntaxError, TemplateNotFound, \
      TemplatesNotFound, TemplateRuntimeError
 from jinja2.utils import import_string, LRUCache, Markup, missing, \
@@ -238,6 +238,14 @@ 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
+
+    #: the context class thatis used for templates.  See
+    #: :class:`~jinja2.runtime.Context` for more information.
+    context_class = Context
+
     def __init__(self,
                  block_start_string=BLOCK_START_STRING,
                  block_end_string=BLOCK_END_STRING,

jinja2/runtime.py

@@ -69,7 +69,8 @@ def new_context(environment, template_name, blocks, vars=None,
         for key, value in iteritems(locals):
             if key[:2] == 'l_' and value is not missing:
                 parent[key[2:]] = value
-    return Context(environment, parent, template_name, blocks)
+    return environment.context_class(environment, parent, template_name,
+                                     blocks)
 
 
 class TemplateReference(object):

tests/test_api.py

@@ -16,6 +16,8 @@ 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.runtime import Context
 from jinja2.utils import Cycler
 
 
@@ -290,3 +292,36 @@ 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'
+
+    def test_custom_context(self):
+        class CustomContext(Context):
+            def resolve(self, key):
+                return 'resolve-' + key
+
+        class CustomEnvironment(Environment):
+            context_class = CustomContext
+
+        env = CustomEnvironment()
+        tmpl = env.from_string('{{ foo }}')
+        assert tmpl.render() == 'resolve-foo'