1from mako import ast
2from mako import exceptions
3from mako import pyparser
4from mako.testing.assertions import assert_raises
5from mako.testing.assertions import eq_
6
7exception_kwargs = {"source": "", "lineno": 0, "pos": 0, "filename": ""}
8
9
10class AstParseTest:
11    def test_locate_identifiers(self):
12        """test the location of identifiers in a python code string"""
13        code = """
14a = 10
15b = 5
16c = x * 5 + a + b + q
17(g,h,i) = (1,2,3)
18[u,k,j] = [4,5,6]
19foo.hoho.lala.bar = 7 + gah.blah + u + blah
20for lar in (1,2,3):
21    gh = 5
22    x = 12
23("hello world, ", a, b)
24("Another expr", c)
25"""
26        parsed = ast.PythonCode(code, **exception_kwargs)
27        eq_(
28            parsed.declared_identifiers,
29            {"a", "b", "c", "g", "h", "i", "u", "k", "j", "gh", "lar", "x"},
30        )
31        eq_(
32            parsed.undeclared_identifiers,
33            {"x", "q", "foo", "gah", "blah"},
34        )
35
36        parsed = ast.PythonCode("x + 5 * (y-z)", **exception_kwargs)
37        assert parsed.undeclared_identifiers == {"x", "y", "z"}
38        assert parsed.declared_identifiers == set()
39
40    def test_locate_identifiers_2(self):
41        code = """
42import foobar
43from lala import hoho, yaya
44import bleep as foo
45result = []
46data = get_data()
47for x in data:
48    result.append(x+7)
49"""
50        parsed = ast.PythonCode(code, **exception_kwargs)
51        eq_(parsed.undeclared_identifiers, {"get_data"})
52        eq_(
53            parsed.declared_identifiers,
54            {"result", "data", "x", "hoho", "foobar", "foo", "yaya"},
55        )
56
57    def test_locate_identifiers_3(self):
58        """test that combination assignment/expressions
59        of the same identifier log the ident as 'undeclared'"""
60        code = """
61x = x + 5
62for y in range(1, y):
63    ("hi",)
64[z for z in range(1, z)]
65(q for q in range (1, q))
66"""
67        parsed = ast.PythonCode(code, **exception_kwargs)
68        eq_(parsed.undeclared_identifiers, {"x", "y", "z", "q", "range"})
69
70    def test_locate_identifiers_4(self):
71        code = """
72x = 5
73(y, )
74def mydef(mydefarg):
75    print("mda is", mydefarg)
76"""
77        parsed = ast.PythonCode(code, **exception_kwargs)
78        eq_(parsed.undeclared_identifiers, {"y"})
79        eq_(parsed.declared_identifiers, {"mydef", "x"})
80
81    def test_locate_identifiers_5(self):
82        code = """
83try:
84    print(x)
85except:
86    print(y)
87"""
88        parsed = ast.PythonCode(code, **exception_kwargs)
89        eq_(parsed.undeclared_identifiers, {"x", "y"})
90
91    def test_locate_identifiers_6(self):
92        code = """
93def foo():
94    return bar()
95"""
96        parsed = ast.PythonCode(code, **exception_kwargs)
97        eq_(parsed.undeclared_identifiers, {"bar"})
98
99        code = """
100def lala(x, y):
101    return x, y, z
102print(x)
103"""
104        parsed = ast.PythonCode(code, **exception_kwargs)
105        eq_(parsed.undeclared_identifiers, {"z", "x"})
106        eq_(parsed.declared_identifiers, {"lala"})
107
108        code = """
109def lala(x, y):
110    def hoho():
111        def bar():
112            z = 7
113print(z)
114"""
115        parsed = ast.PythonCode(code, **exception_kwargs)
116        eq_(parsed.undeclared_identifiers, {"z"})
117        eq_(parsed.declared_identifiers, {"lala"})
118
119    def test_locate_identifiers_7(self):
120        code = """
121import foo.bar
122"""
123        parsed = ast.PythonCode(code, **exception_kwargs)
124        eq_(parsed.declared_identifiers, {"foo"})
125        eq_(parsed.undeclared_identifiers, set())
126
127    def test_locate_identifiers_8(self):
128        code = """
129class Hi:
130    foo = 7
131    def hoho(self):
132        x = 5
133"""
134        parsed = ast.PythonCode(code, **exception_kwargs)
135        eq_(parsed.declared_identifiers, {"Hi"})
136        eq_(parsed.undeclared_identifiers, set())
137
138    def test_locate_identifiers_9(self):
139        code = """
140    ",".join([t for t in ("a", "b", "c")])
141"""
142        parsed = ast.PythonCode(code, **exception_kwargs)
143        eq_(parsed.declared_identifiers, {"t"})
144        eq_(parsed.undeclared_identifiers, {"t"})
145
146        code = """
147    [(val, name) for val, name in x]
148"""
149        parsed = ast.PythonCode(code, **exception_kwargs)
150        eq_(parsed.declared_identifiers, {"val", "name"})
151        eq_(parsed.undeclared_identifiers, {"val", "name", "x"})
152
153    def test_locate_identifiers_10(self):
154        code = """
155lambda q: q + 5
156"""
157        parsed = ast.PythonCode(code, **exception_kwargs)
158        eq_(parsed.declared_identifiers, set())
159        eq_(parsed.undeclared_identifiers, set())
160
161    def test_locate_identifiers_11(self):
162        code = """
163def x(q):
164    return q + 5
165"""
166        parsed = ast.PythonCode(code, **exception_kwargs)
167        eq_(parsed.declared_identifiers, {"x"})
168        eq_(parsed.undeclared_identifiers, set())
169
170    def test_locate_identifiers_12(self):
171        code = """
172def foo():
173    s = 1
174    def bar():
175        t = s
176"""
177        parsed = ast.PythonCode(code, **exception_kwargs)
178        eq_(parsed.declared_identifiers, {"foo"})
179        eq_(parsed.undeclared_identifiers, set())
180
181    def test_locate_identifiers_13(self):
182        code = """
183def foo():
184    class Bat:
185        pass
186    Bat
187"""
188        parsed = ast.PythonCode(code, **exception_kwargs)
189        eq_(parsed.declared_identifiers, {"foo"})
190        eq_(parsed.undeclared_identifiers, set())
191
192    def test_locate_identifiers_14(self):
193        code = """
194def foo():
195    class Bat:
196        pass
197    Bat
198
199print(Bat)
200"""
201        parsed = ast.PythonCode(code, **exception_kwargs)
202        eq_(parsed.declared_identifiers, {"foo"})
203        eq_(parsed.undeclared_identifiers, {"Bat"})
204
205    def test_locate_identifiers_16(self):
206        code = """
207try:
208    print(x)
209except Exception as e:
210    print(y)
211"""
212        parsed = ast.PythonCode(code, **exception_kwargs)
213        eq_(parsed.undeclared_identifiers, {"x", "y", "Exception"})
214
215    def test_locate_identifiers_17(self):
216        code = """
217try:
218    print(x)
219except (Foo, Bar) as e:
220    print(y)
221"""
222        parsed = ast.PythonCode(code, **exception_kwargs)
223        eq_(parsed.undeclared_identifiers, {"x", "y", "Foo", "Bar"})
224
225    def test_no_global_imports(self):
226        code = """
227from foo import *
228import x as bar
229"""
230        assert_raises(
231            exceptions.CompileException,
232            ast.PythonCode,
233            code,
234            **exception_kwargs,
235        )
236
237    def test_python_fragment(self):
238        parsed = ast.PythonFragment("for x in foo:", **exception_kwargs)
239        eq_(parsed.declared_identifiers, {"x"})
240        eq_(parsed.undeclared_identifiers, {"foo"})
241
242        parsed = ast.PythonFragment("try:", **exception_kwargs)
243
244        parsed = ast.PythonFragment(
245            "except MyException as e:", **exception_kwargs
246        )
247        eq_(parsed.declared_identifiers, {"e"})
248        eq_(parsed.undeclared_identifiers, {"MyException"})
249
250    def test_argument_list(self):
251        parsed = ast.ArgumentList(
252            "3, 5, 'hi', x+5, " "context.get('lala')", **exception_kwargs
253        )
254        eq_(parsed.undeclared_identifiers, {"x", "context"})
255        eq_(
256            [x for x in parsed.args],
257            ["3", "5", "'hi'", "(x + 5)", "context.get('lala')"],
258        )
259
260        parsed = ast.ArgumentList("h", **exception_kwargs)
261        eq_(parsed.args, ["h"])
262
263    def test_function_decl(self):
264        """test getting the arguments from a function"""
265        code = "def foo(a, b, c=None, d='hi', e=x, f=y+7):pass"
266        parsed = ast.FunctionDecl(code, **exception_kwargs)
267        eq_(parsed.funcname, "foo")
268        eq_(parsed.argnames, ["a", "b", "c", "d", "e", "f"])
269        eq_(parsed.kwargnames, [])
270
271    def test_function_decl_2(self):
272        """test getting the arguments from a function"""
273        code = "def foo(a, b, c=None, *args, **kwargs):pass"
274        parsed = ast.FunctionDecl(code, **exception_kwargs)
275        eq_(parsed.funcname, "foo")
276        eq_(parsed.argnames, ["a", "b", "c", "args"])
277        eq_(parsed.kwargnames, ["kwargs"])
278
279    def test_function_decl_3(self):
280        """test getting the arguments from a fancy py3k function"""
281        code = "def foo(a, b, *c, d, e, **f):pass"
282        parsed = ast.FunctionDecl(code, **exception_kwargs)
283        eq_(parsed.funcname, "foo")
284        eq_(parsed.argnames, ["a", "b", "c"])
285        eq_(parsed.kwargnames, ["d", "e", "f"])
286
287    def test_expr_generate(self):
288        """test the round trip of expressions to AST back to python source"""
289        x = 1
290        y = 2
291
292        class F:
293            def bar(self, a, b):
294                return a + b
295
296        def lala(arg):
297            return "blah" + arg
298
299        local_dict = dict(x=x, y=y, foo=F(), lala=lala)
300
301        code = "str((x+7*y) / foo.bar(5,6)) + lala('ho')"
302        astnode = pyparser.parse(code)
303        newcode = pyparser.ExpressionGenerator(astnode).value()
304        eq_(eval(code, local_dict), eval(newcode, local_dict))
305
306        a = ["one", "two", "three"]
307        hoho = {"somevalue": "asdf"}
308        g = [1, 2, 3, 4, 5]
309        local_dict = dict(a=a, hoho=hoho, g=g)
310        code = (
311            "a[2] + hoho['somevalue'] + "
312            "repr(g[3:5]) + repr(g[3:]) + repr(g[:5])"
313        )
314        astnode = pyparser.parse(code)
315        newcode = pyparser.ExpressionGenerator(astnode).value()
316        eq_(eval(code, local_dict), eval(newcode, local_dict))
317
318        local_dict = {"f": lambda: 9, "x": 7}
319        code = "x+f()"
320        astnode = pyparser.parse(code)
321        newcode = pyparser.ExpressionGenerator(astnode).value()
322        eq_(eval(code, local_dict), eval(newcode, local_dict))
323
324        for code in [
325            "repr({'x':7,'y':18})",
326            "repr([])",
327            "repr({})",
328            "repr([{3:[]}])",
329            "repr({'x':37*2 + len([6,7,8])})",
330            "repr([1, 2, {}, {'x':'7'}])",
331            "repr({'x':-1})",
332            "repr(((1,2,3), (4,5,6)))",
333            "repr(1 and 2 and 3 and 4)",
334            "repr(True and False or 55)",
335            "repr(lambda x, y: (x + y))",
336            "repr(lambda *arg, **kw: arg, kw)",
337            "repr(1 & 2 | 3)",
338            "repr(3//5)",
339            "repr(3^5)",
340            "repr([q.endswith('e') for q in " "['one', 'two', 'three']])",
341            "repr([x for x in (5,6,7) if x == 6])",
342            "repr(not False)",
343        ]:
344            local_dict = {}
345            astnode = pyparser.parse(code)
346            newcode = pyparser.ExpressionGenerator(astnode).value()
347            if "lambda" in code:
348                eq_(code, newcode)
349            else:
350                eq_(eval(code, local_dict), eval(newcode, local_dict))
351