track local loop/block vars for contextfunctions

4 files changed, 148 insertions, 4 deletions

diff --git a/CHANGES.rst b/CHANGES.rst
index c53d67e..3ea51cb 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -32,6 +32,10 @@ Unreleased
     :issue:`522, 827, 1172`, :pr:`1195`
 -   Filters that get attributes, such as ``map`` and ``groupby``, can
     use a false or empty value as a default. :issue:`1331`
+-   Fix a bug that prevented variables set in blocks or loops from
+    being accessed in custom context functions. :issue:`768`
+-   Fix a bug that caused scoped blocks from accessing special loop
+    variables. :issue:`1088`
 
 
 Version 2.11.3
diff --git a/src/jinja2/compiler.py b/src/jinja2/compiler.py
index 67918cd..50f98f1 100644
--- a/src/jinja2/compiler.py
+++ b/src/jinja2/compiler.py
@@ -131,6 +131,11 @@ class Frame:
         if parent is not None:
             self.buffer = parent.buffer
 
+        # variables set inside of loops and blocks should not affect outer frames,
+        # but they still needs to be kept track of as part of the active context.
+        self.loop_frame = False
+        self.block_frame = False
+
     def copy(self):
         """Create a copy of the current one."""
         rv = object.__new__(self.__class__)
@@ -639,22 +644,38 @@ class CodeGenerator(NodeVisitor):
         context variables if necessary.
         """
         vars = self._assign_stack.pop()
-        if not frame.toplevel or not vars:
+        if (
+            not frame.block_frame
+            and not frame.loop_frame
+            and not frame.toplevel
+            or not vars
+        ):
             return
         public_names = [x for x in vars if x[:1] != "_"]
         if len(vars) == 1:
             name = next(iter(vars))
             ref = frame.symbols.ref(name)
+            if frame.loop_frame:
+                self.writeline(f"_loop_vars[{name!r}] = {ref}")
+                return
+            if frame.block_frame:
+                self.writeline(f"_block_vars[{name!r}] = {ref}")
+                return
             self.writeline(f"context.vars[{name!r}] = {ref}")
         else:
-            self.writeline("context.vars.update({")
+            if frame.loop_frame:
+                self.writeline("_loop_vars.update({")
+            elif frame.block_frame:
+                self.writeline("_block_vars.update({")
+            else:
+                self.writeline("context.vars.update({")
             for idx, name in enumerate(vars):
                 if idx:
                     self.write(", ")
                 ref = frame.symbols.ref(name)
                 self.write(f"{name!r}: {ref}")
             self.write("})")
-        if public_names:
+        if not frame.block_frame and not frame.loop_frame and public_names:
             if len(public_names) == 1:
                 self.writeline(f"context.exported_vars.add({public_names[0]!r})")
             else:
@@ -760,6 +781,7 @@ class CodeGenerator(NodeVisitor):
             # toplevel template.  This would cause a variety of
             # interesting issues with identifier tracking.
             block_frame = Frame(eval_ctx)
+            block_frame.block_frame = True
             undeclared = find_undeclared(block.body, ("self", "super"))
             if "self" in undeclared:
                 ref = block_frame.symbols.declare_parameter("self")
@@ -769,6 +791,7 @@ class CodeGenerator(NodeVisitor):
                 self.writeline(f"{ref} = context.super({name!r}, block_{name})")
             block_frame.symbols.analyze_node(block)
             block_frame.block = name
+            self.writeline("_block_vars = {}")
             self.enter_frame(block_frame)
             self.pull_dependencies(block.body)
             self.blockvisit(block.body, block_frame)
@@ -1003,6 +1026,7 @@ class CodeGenerator(NodeVisitor):
 
     def visit_For(self, node, frame):
         loop_frame = frame.inner()
+        loop_frame.loop_frame = True
         test_frame = frame.inner()
         else_frame = frame.inner()
 
@@ -1103,6 +1127,7 @@ class CodeGenerator(NodeVisitor):
         self.indent()
         self.enter_frame(loop_frame)
 
+        self.writeline("_loop_vars = {}")
         self.blockvisit(node.body, loop_frame)
         if node.else_:
             self.writeline(f"{iteration_indicator} = 0")
@@ -1411,7 +1436,9 @@ class CodeGenerator(NodeVisitor):
     # -- Expression Visitors
 
     def visit_Name(self, node, frame):
-        if node.ctx == "store" and frame.toplevel:
+        if node.ctx == "store" and (
+            frame.toplevel or frame.loop_frame or frame.block_frame
+        ):
             if self._assign_stack:
                 self._assign_stack[-1].add(node.name)
         ref = frame.symbols.ref(node.name)
@@ -1679,6 +1706,12 @@ class CodeGenerator(NodeVisitor):
             self.write("context.call(")
         self.visit(node.node, frame)
         extra_kwargs = {"caller": "caller"} if forward_caller else None
+        loop_kwargs = {"_loop_vars": "_loop_vars"} if frame.loop_frame else {}
+        block_kwargs = {"_block_vars": "_block_vars"} if frame.block_frame else {}
+        if extra_kwargs:
+            extra_kwargs.update(loop_kwargs, **block_kwargs)
+        elif loop_kwargs or block_kwargs:
+            extra_kwargs = dict(loop_kwargs, **block_kwargs)
         self.signature(node, frame, extra_kwargs)
         self.write(")")
         if self.environment.is_async:
diff --git a/src/jinja2/runtime.py b/src/jinja2/runtime.py
index a05b196..4a3c36e 100644
--- a/src/jinja2/runtime.py
+++ b/src/jinja2/runtime.py
@@ -284,11 +284,20 @@ class Context(metaclass=ContextMeta):
 
         if callable(__obj):
             if getattr(__obj, "contextfunction", False) is True:
+                # the active context should have access to variables set in
+                # loops and blocks without mutating the context itself
+                if kwargs.get("_loop_vars"):
+                    __self = __self.derived(kwargs["_loop_vars"])
+                if kwargs.get("_block_vars"):
+                    __self = __self.derived(kwargs["_block_vars"])
                 args = (__self,) + args
             elif getattr(__obj, "evalcontextfunction", False) is True:
                 args = (__self.eval_ctx,) + args
             elif getattr(__obj, "environmentfunction", False) is True:
                 args = (__self.environment,) + args
+
+        kwargs.pop("_block_vars", None)
+        kwargs.pop("_loop_vars", None)
         try:
             return __obj(*args, **kwargs)
         except StopIteration:
diff --git a/tests/test_regression.py b/tests/test_regression.py
index 21a6d92..716d4a0 100644
--- a/tests/test_regression.py
+++ b/tests/test_regression.py
@@ -7,6 +7,7 @@ from jinja2 import Template
 from jinja2 import TemplateAssertionError
 from jinja2 import TemplateNotFound
 from jinja2 import TemplateSyntaxError
+from jinja2.utils import contextfunction
 
 
 class TestCorner:
@@ -618,3 +619,100 @@ class TestBug:
         from jinja2.runtime import ChainableUndefined
 
         assert str(Markup(ChainableUndefined())) == ""
+
+    def test_scoped_block_loop_vars(self, env):
+        tmpl = env.from_string(
+            """\
+Start
+{% for i in ["foo", "bar"] -%}
+{% block body scoped -%}
+{{ loop.index }}) {{ i }}{% if loop.last %} last{% endif -%}
+{%- endblock %}
+{% endfor -%}
+End"""
+        )
+        assert tmpl.render() == "Start\n1) foo\n2) bar last\nEnd"
+
+    def test_contextfunction_loop_vars(self, env):
+        @contextfunction
+        def test(ctx):
+            return f"{ctx['i']}{ctx['j']}"
+
+        tmpl = env.from_string(
+            """\
+{% set i = 42 %}
+{%- for idx in range(2) -%}
+{{ i }}{{ j }}
+{% set i = idx -%}
+{%- set j = loop.index -%}
+{{ test() }}
+{{ i }}{{ j }}
+{% endfor -%}
+{{ i }}{{ j }}"""
+        )
+        tmpl.globals["test"] = test
+        assert tmpl.render() == "42\n01\n01\n42\n12\n12\n42"
+
+    def test_contextfunction_scoped_loop_vars(self, env):
+        @contextfunction
+        def test(ctx):
+            return f"{ctx['i']}"
+
+        tmpl = env.from_string(
+            """\
+{% set i = 42 %}
+{%- for idx in range(2) -%}
+{{ i }}
+{%- set i = loop.index0 -%}
+{% block body scoped %}
+{{ test() }}
+{% endblock -%}
+{% endfor -%}
+{{ i }}"""
+        )
+        tmpl.globals["test"] = test
+        assert tmpl.render() == "42\n0\n42\n1\n42"
+
+    def test_contextfunction_in_blocks(self, env):
+        @contextfunction
+        def test(ctx):
+            return f"{ctx['i']}"
+
+        tmpl = env.from_string(
+            """\
+{%- set i = 42 -%}
+{{ i }}
+{% block body -%}
+{% set i = 24 -%}
+{{ test() }}
+{% endblock -%}
+{{ i }}"""
+        )
+        tmpl.globals["test"] = test
+        assert tmpl.render() == "42\n24\n42"
+
+    def test_contextfunction_block_and_loop(self, env):
+        @contextfunction
+        def test(ctx):
+            return f"{ctx['i']}"
+
+        tmpl = env.from_string(
+            """\
+{%- set i = 42 -%}
+{% for idx in range(2) -%}
+{{ test() }}
+{%- set i = idx -%}
+{% block body scoped %}
+{{ test() }}
+{% set i = 24 -%}
+{{ test() }}
+{% endblock -%}
+{{ test() }}
+{% endfor -%}
+{{ test() }}"""
+        )
+        tmpl.globals["test"] = test
+
+        # values set within a block or loop should not
+        # show up outside of it
+        assert tmpl.render() == "42\n0\n24\n0\n42\n1\n24\n1\n42"