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