track local loop/block vars for contextfunctions

This commit is contained in:
Amy 2020-06-23 10:53:59 -04:00
parent fed1b24d5f
commit f524bcce0c
4 changed files with 148 additions and 4 deletions

View File

@ -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

View File

@ -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:

View File

@ -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:

View File

@ -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"