1# mako/ast.py
2# Copyright (C) 2006-2015 the Mako authors and contributors <see AUTHORS file>
3#
4# This module is part of Mako and is released under
5# the MIT License: http://www.opensource.org/licenses/mit-license.php
6
7"""utilities for analyzing expressions and blocks of Python
8code, as well as generating Python from AST nodes"""
9
10from mako import exceptions, pyparser, compat
11import re
12
13class PythonCode(object):
14    """represents information about a string containing Python code"""
15    def __init__(self, code, **exception_kwargs):
16        self.code = code
17
18        # represents all identifiers which are assigned to at some point in
19        # the code
20        self.declared_identifiers = set()
21
22        # represents all identifiers which are referenced before their
23        # assignment, if any
24        self.undeclared_identifiers = set()
25
26        # note that an identifier can be in both the undeclared and declared
27        # lists.
28
29        # using AST to parse instead of using code.co_varnames,
30        # code.co_names has several advantages:
31        # - we can locate an identifier as "undeclared" even if
32        # its declared later in the same block of code
33        # - AST is less likely to break with version changes
34        # (for example, the behavior of co_names changed a little bit
35        # in python version 2.5)
36        if isinstance(code, compat.string_types):
37            expr = pyparser.parse(code.lstrip(), "exec", **exception_kwargs)
38        else:
39            expr = code
40
41        f = pyparser.FindIdentifiers(self, **exception_kwargs)
42        f.visit(expr)
43
44class ArgumentList(object):
45    """parses a fragment of code as a comma-separated list of expressions"""
46    def __init__(self, code, **exception_kwargs):
47        self.codeargs = []
48        self.args = []
49        self.declared_identifiers = set()
50        self.undeclared_identifiers = set()
51        if isinstance(code, compat.string_types):
52            if re.match(r"\S", code) and not re.match(r",\s*$", code):
53                # if theres text and no trailing comma, insure its parsed
54                # as a tuple by adding a trailing comma
55                code  += ","
56            expr = pyparser.parse(code, "exec", **exception_kwargs)
57        else:
58            expr = code
59
60        f = pyparser.FindTuple(self, PythonCode, **exception_kwargs)
61        f.visit(expr)
62
63class PythonFragment(PythonCode):
64    """extends PythonCode to provide identifier lookups in partial control
65    statements
66
67    e.g.
68        for x in 5:
69        elif y==9:
70        except (MyException, e):
71    etc.
72    """
73    def __init__(self, code, **exception_kwargs):
74        m = re.match(r'^(\w+)(?:\s+(.*?))?:\s*(#|$)', code.strip(), re.S)
75        if not m:
76            raise exceptions.CompileException(
77                          "Fragment '%s' is not a partial control statement" %
78                          code, **exception_kwargs)
79        if m.group(3):
80            code = code[:m.start(3)]
81        (keyword, expr) = m.group(1,2)
82        if keyword in ['for','if', 'while']:
83            code = code + "pass"
84        elif keyword == 'try':
85            code = code + "pass\nexcept:pass"
86        elif keyword == 'elif' or keyword == 'else':
87            code = "if False:pass\n" + code + "pass"
88        elif keyword == 'except':
89            code = "try:pass\n" + code + "pass"
90        elif keyword == 'with':
91            code = code + "pass"
92        else:
93            raise exceptions.CompileException(
94                                "Unsupported control keyword: '%s'" %
95                                keyword, **exception_kwargs)
96        super(PythonFragment, self).__init__(code, **exception_kwargs)
97
98
99class FunctionDecl(object):
100    """function declaration"""
101    def __init__(self, code, allow_kwargs=True, **exception_kwargs):
102        self.code = code
103        expr = pyparser.parse(code, "exec", **exception_kwargs)
104
105        f = pyparser.ParseFunc(self, **exception_kwargs)
106        f.visit(expr)
107        if not hasattr(self, 'funcname'):
108            raise exceptions.CompileException(
109                            "Code '%s' is not a function declaration" % code,
110                            **exception_kwargs)
111        if not allow_kwargs and self.kwargs:
112            raise exceptions.CompileException(
113                                "'**%s' keyword argument not allowed here" %
114                                self.kwargnames[-1], **exception_kwargs)
115
116    def get_argument_expressions(self, as_call=False):
117        """Return the argument declarations of this FunctionDecl as a printable
118        list.
119
120        By default the return value is appropriate for writing in a ``def``;
121        set `as_call` to true to build arguments to be passed to the function
122        instead (assuming locals with the same names as the arguments exist).
123        """
124
125        namedecls = []
126
127        # Build in reverse order, since defaults and slurpy args come last
128        argnames = self.argnames[::-1]
129        kwargnames = self.kwargnames[::-1]
130        defaults = self.defaults[::-1]
131        kwdefaults = self.kwdefaults[::-1]
132
133        # Named arguments
134        if self.kwargs:
135            namedecls.append("**" + kwargnames.pop(0))
136
137        for name in kwargnames:
138            # Keyword-only arguments must always be used by name, so even if
139            # this is a call, print out `foo=foo`
140            if as_call:
141                namedecls.append("%s=%s" % (name, name))
142            elif kwdefaults:
143                default = kwdefaults.pop(0)
144                if default is None:
145                    # The AST always gives kwargs a default, since you can do
146                    # `def foo(*, a=1, b, c=3)`
147                    namedecls.append(name)
148                else:
149                    namedecls.append("%s=%s" % (
150                        name, pyparser.ExpressionGenerator(default).value()))
151            else:
152                namedecls.append(name)
153
154        # Positional arguments
155        if self.varargs:
156            namedecls.append("*" + argnames.pop(0))
157
158        for name in argnames:
159            if as_call or not defaults:
160                namedecls.append(name)
161            else:
162                default = defaults.pop(0)
163                namedecls.append("%s=%s" % (
164                    name, pyparser.ExpressionGenerator(default).value()))
165
166        namedecls.reverse()
167        return namedecls
168
169    @property
170    def allargnames(self):
171        return tuple(self.argnames) + tuple(self.kwargnames)
172
173class FunctionArgs(FunctionDecl):
174    """the argument portion of a function declaration"""
175
176    def __init__(self, code, **kwargs):
177        super(FunctionArgs, self).__init__("def ANON(%s):pass" % code,
178                **kwargs)
179