Merge pull request #861 from OddBloke/contextreference_locals

Add a DerivedContextReference node type

4 files changed, 36 insertions, 4 deletions

diff --git a/CHANGES.rst b/CHANGES.rst
index 037dfaa..f5f2e7c 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -85,6 +85,9 @@ Unreleased
     and source for Python >= 3.7. :issue:`1104`
 -   Tracebacks for template syntax errors in Python 3 no longer show
     internal compiler frames. :issue:`763`
+-   Add a ``DerivedContextReference`` node that can be used by
+    extensions to get the current context and local variables such as
+    ``loop``. :issue:`860`
 
 
 Version 2.10.3
diff --git a/jinja2/compiler.py b/jinja2/compiler.py
index 3c69df8..addf71e 100644
--- a/jinja2/compiler.py
+++ b/jinja2/compiler.py
@@ -1720,6 +1720,9 @@ class CodeGenerator(NodeVisitor):
     def visit_ContextReference(self, node, frame):
         self.write('context')
 
+    def visit_DerivedContextReference(self, node, frame):
+        self.write(self.derive_context(frame))
+
     def visit_Continue(self, node, frame):
         self.writeline('continue', node)
 
diff --git a/jinja2/nodes.py b/jinja2/nodes.py
index e46aa91..00e0329 100644
--- a/jinja2/nodes.py
+++ b/jinja2/nodes.py
@@ -952,6 +952,15 @@ class ContextReference(Expr):
     """
 
 
+class DerivedContextReference(Expr):
+    """Return the current template context including locals. Behaves
+    exactly like :class:`ContextReference`, but includes local
+    variables, such as from a ``for`` loop.
+
+    .. versionadded:: 2.11
+    """
+
+
 class Continue(Stmt):
     """Continue a loop."""
 
diff --git a/tests/test_ext.py b/tests/test_ext.py
index 1c349bd..7975271 100644
--- a/tests/test_ext.py
+++ b/tests/test_ext.py
@@ -109,24 +109,30 @@ newstyle_i18n_env.install_gettext_callables(gettext, ngettext, newstyle=True)
 class ExampleExtension(Extension):
     tags = set(['test'])
     ext_attr = 42
+    context_reference_node_cls = nodes.ContextReference
 
     def parse(self, parser):
         return nodes.Output([self.call_method('_dump', [
             nodes.EnvironmentAttribute('sandboxed'),
             self.attr('ext_attr'),
             nodes.ImportedName(__name__ + '.importable_object'),
-            nodes.ContextReference()
+            self.context_reference_node_cls()
         ])]).set_lineno(next(parser.stream).lineno)
 
     def _dump(self, sandboxed, ext_attr, imported_object, context):
-        return '%s|%s|%s|%s' % (
+        return '%s|%s|%s|%s|%s' % (
             sandboxed,
             ext_attr,
             imported_object,
-            context.blocks
+            context.blocks,
+            context.get('test_var')
         )
 
 
+class DerivedExampleExtension(ExampleExtension):
+    context_reference_node_cls = nodes.DerivedContextReference
+
+
 class PreprocessorExtension(Extension):
 
     def preprocess(self, source, name, filename=None):
@@ -205,7 +211,18 @@ class TestExtensions(object):
     def test_extension_nodes(self):
         env = Environment(extensions=[ExampleExtension])
         tmpl = env.from_string('{% test %}')
-        assert tmpl.render() == 'False|42|23|{}'
+        assert tmpl.render() == 'False|42|23|{}|None'
+
+    def test_contextreference_node_passes_context(self):
+        env = Environment(extensions=[ExampleExtension])
+        tmpl = env.from_string('{% set test_var="test_content" %}{% test %}')
+        assert tmpl.render() == 'False|42|23|{}|test_content'
+
+    def test_contextreference_node_can_pass_locals(self):
+        env = Environment(extensions=[DerivedExampleExtension])
+        tmpl = env.from_string(
+            '{% for test_var in ["test_content"] %}{% test %}{% endfor %}')
+        assert tmpl.render() == 'False|42|23|{}|test_content'
 
     def test_identifier(self):
         assert ExampleExtension.identifier == __name__ + '.ExampleExtension'