Skip to content

Commit e2606d5

Browse files
cocolatosqla-tester
authored andcommitted
Support comprehensions inside functions when use strict_undefined flag.
Fixes: #320 Now the test code works as expected if strict_undefined is set to true: ```python from mako.template import Template text = """ <% mydict = { 'foo': 1 } ## Uncomment the following line to workaround the error ##k = None def getkeys(x): return [ k for k in x.keys() ] %> ${ ','.join( getkeys(mydict) ) } """ tmpl = Template(text=text, strict_undefined=True) out = tmpl.render() print(out) ``` output: ``` foo ``` Closes: #386 Pull-request: #386 Pull-request-sha: cc6a3e0 Change-Id: I0591873a83837f8f35b0963c0536df1e2675012f
1 parent 2815589 commit e2606d5

File tree

4 files changed

+124
-0
lines changed

4 files changed

+124
-0
lines changed

doc/build/unreleased/320.rst

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
.. change::
2+
:tags: bug, parser
3+
:tickets: 320
4+
5+
Fixed unexpected syntax error in strict_undefined mode that occurred
6+
when using comprehensions within a function in a Mako Python code block.
7+
Now, the local variable in comprehensions won't be added to the checklist
8+
when using strict_undefined mode.
9+
Pull request courtesy Hai Zhu.

mako/pyparser.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,26 @@ def visit_FunctionDef(self, node):
9090
self._add_declared(node.name)
9191
self._visit_function(node, False)
9292

93+
def visit_ListComp(self, node):
94+
if self.in_function:
95+
if not isinstance(node.elt, _ast.Name):
96+
self.visit(node.elt)
97+
for comp in node.generators:
98+
self.visit(comp.iter)
99+
else:
100+
self.generic_visit(node)
101+
102+
visit_SetComp = visit_GeneratorExp = visit_ListComp
103+
104+
def visit_DictComp(self, node):
105+
if self.in_function:
106+
if not isinstance(node.key, _ast.Name):
107+
self.visit(node.elt)
108+
for comp in node.generators:
109+
self.visit(comp.iter)
110+
else:
111+
self.generic_visit(node)
112+
93113
def _expand_tuples(self, args):
94114
for arg in args:
95115
if isinstance(arg, _ast.Tuple):

test/test_ast.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -222,6 +222,42 @@ def test_locate_identifiers_17(self):
222222
parsed = ast.PythonCode(code, **exception_kwargs)
223223
eq_(parsed.undeclared_identifiers, {"x", "y", "Foo", "Bar"})
224224

225+
def test_locate_identifiers_18(self):
226+
code = """
227+
def func():
228+
return [i for i in range(10)]
229+
"""
230+
parsed = ast.PythonCode(code, **exception_kwargs)
231+
eq_(parsed.declared_identifiers, {"func"})
232+
eq_(parsed.undeclared_identifiers, {"range"})
233+
234+
def test_locate_identifiers_19(self):
235+
code = """
236+
def func():
237+
return (i for i in range(10))
238+
"""
239+
parsed = ast.PythonCode(code, **exception_kwargs)
240+
eq_(parsed.declared_identifiers, {"func"})
241+
eq_(parsed.undeclared_identifiers, {"range"})
242+
243+
def test_locate_identifiers_20(self):
244+
code = """
245+
def func():
246+
return {i for i in range(10)}
247+
"""
248+
parsed = ast.PythonCode(code, **exception_kwargs)
249+
eq_(parsed.declared_identifiers, {"func"})
250+
eq_(parsed.undeclared_identifiers, {"range"})
251+
252+
def test_locate_identifiers_21(self):
253+
code = """
254+
def func():
255+
return {i: i**2 for i in range(10)}
256+
"""
257+
parsed = ast.PythonCode(code, **exception_kwargs)
258+
eq_(parsed.declared_identifiers, {"func"})
259+
eq_(parsed.undeclared_identifiers, {"range"})
260+
225261
def test_no_global_imports(self):
226262
code = """
227263
from foo import *

test/test_template.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1717,3 +1717,62 @@ def test_inline_percent(self):
17171717
"% foo",
17181718
"bar %% baz",
17191719
]
1720+
1721+
def test_listcomp_in_func_strict(self):
1722+
t = Template(
1723+
"""
1724+
<%
1725+
mydict = { 'foo': 1 }
1726+
def getkeys(x):
1727+
return [ k for k in x.keys() ]
1728+
%>
1729+
1730+
${ ','.join( getkeys(mydict) ) }
1731+
""",
1732+
strict_undefined=True,
1733+
)
1734+
assert result_raw_lines(t.render()) == ["foo"]
1735+
1736+
def test_setcomp_in_func_strict(self):
1737+
t = Template(
1738+
"""
1739+
<%
1740+
mydict = { 'foo': 1 }
1741+
def getkeys(x):
1742+
return { k for k in x.keys() }
1743+
%>
1744+
1745+
${ ','.join( getkeys(mydict) ) }
1746+
""",
1747+
strict_undefined=True,
1748+
)
1749+
assert result_raw_lines(t.render()) == ["foo"]
1750+
1751+
def test_generator_in_func_strict(self):
1752+
t = Template(
1753+
"""
1754+
<%
1755+
mydict = { 'foo': 1 }
1756+
def getkeys(x):
1757+
return ( k for k in x.keys())
1758+
%>
1759+
1760+
${ ','.join( getkeys(mydict) ) }
1761+
""",
1762+
strict_undefined=True,
1763+
)
1764+
assert result_raw_lines(t.render()) == ["foo"]
1765+
1766+
def test_dictcomp_in_func_strict(self):
1767+
t = Template(
1768+
"""
1769+
<%
1770+
def square():
1771+
return {i: i**2 for i in range(10)}
1772+
%>
1773+
1774+
${ square()[3] }
1775+
""",
1776+
strict_undefined=True,
1777+
)
1778+
assert result_raw_lines(t.render()) == ["9"]

0 commit comments

Comments
 (0)