1import platform
2import sys
3import typing as t
4from types import CodeType
5from types import TracebackType
6
7from .exceptions import TemplateSyntaxError
8from .utils import internal_code
9from .utils import missing
10
11if t.TYPE_CHECKING:
12    from .runtime import Context
13
14
15def rewrite_traceback_stack(source: t.Optional[str] = None) -> BaseException:
16    """Rewrite the current exception to replace any tracebacks from
17    within compiled template code with tracebacks that look like they
18    came from the template source.
19
20    This must be called within an ``except`` block.
21
22    :param source: For ``TemplateSyntaxError``, the original source if
23        known.
24    :return: The original exception with the rewritten traceback.
25    """
26    _, exc_value, tb = sys.exc_info()
27    exc_value = t.cast(BaseException, exc_value)
28    tb = t.cast(TracebackType, tb)
29
30    if isinstance(exc_value, TemplateSyntaxError) and not exc_value.translated:
31        exc_value.translated = True
32        exc_value.source = source
33        # Remove the old traceback, otherwise the frames from the
34        # compiler still show up.
35        exc_value.with_traceback(None)
36        # Outside of runtime, so the frame isn't executing template
37        # code, but it still needs to point at the template.
38        tb = fake_traceback(
39            exc_value, None, exc_value.filename or "<unknown>", exc_value.lineno
40        )
41    else:
42        # Skip the frame for the render function.
43        tb = tb.tb_next
44
45    stack = []
46
47    # Build the stack of traceback object, replacing any in template
48    # code with the source file and line information.
49    while tb is not None:
50        # Skip frames decorated with @internalcode. These are internal
51        # calls that aren't useful in template debugging output.
52        if tb.tb_frame.f_code in internal_code:
53            tb = tb.tb_next
54            continue
55
56        template = tb.tb_frame.f_globals.get("__jinja_template__")
57
58        if template is not None:
59            lineno = template.get_corresponding_lineno(tb.tb_lineno)
60            fake_tb = fake_traceback(exc_value, tb, template.filename, lineno)
61            stack.append(fake_tb)
62        else:
63            stack.append(tb)
64
65        tb = tb.tb_next
66
67    tb_next = None
68
69    # Assign tb_next in reverse to avoid circular references.
70    for tb in reversed(stack):
71        tb_next = tb_set_next(tb, tb_next)
72
73    return exc_value.with_traceback(tb_next)
74
75
76def fake_traceback(  # type: ignore
77    exc_value: BaseException, tb: t.Optional[TracebackType], filename: str, lineno: int
78) -> TracebackType:
79    """Produce a new traceback object that looks like it came from the
80    template source instead of the compiled code. The filename, line
81    number, and location name will point to the template, and the local
82    variables will be the current template context.
83
84    :param exc_value: The original exception to be re-raised to create
85        the new traceback.
86    :param tb: The original traceback to get the local variables and
87        code info from.
88    :param filename: The template filename.
89    :param lineno: The line number in the template source.
90    """
91    if tb is not None:
92        # Replace the real locals with the context that would be
93        # available at that point in the template.
94        locals = get_template_locals(tb.tb_frame.f_locals)
95        locals.pop("__jinja_exception__", None)
96    else:
97        locals = {}
98
99    globals = {
100        "__name__": filename,
101        "__file__": filename,
102        "__jinja_exception__": exc_value,
103    }
104    # Raise an exception at the correct line number.
105    code = compile("\n" * (lineno - 1) + "raise __jinja_exception__", filename, "exec")
106
107    # Build a new code object that points to the template file and
108    # replaces the location with a block name.
109    try:
110        location = "template"
111
112        if tb is not None:
113            function = tb.tb_frame.f_code.co_name
114
115            if function == "root":
116                location = "top-level template code"
117            elif function.startswith("block_"):
118                location = f"block {function[6:]!r}"
119
120        # Collect arguments for the new code object. CodeType only
121        # accepts positional arguments, and arguments were inserted in
122        # new Python versions.
123        code_args = []
124
125        for attr in (
126            "argcount",
127            "posonlyargcount",  # Python 3.8
128            "kwonlyargcount",
129            "nlocals",
130            "stacksize",
131            "flags",
132            "code",  # codestring
133            "consts",  # constants
134            "names",
135            "varnames",
136            ("filename", filename),
137            ("name", location),
138            "firstlineno",
139            "lnotab",
140            "freevars",
141            "cellvars",
142            "linetable",  # Python 3.10
143        ):
144            if isinstance(attr, tuple):
145                # Replace with given value.
146                code_args.append(attr[1])
147                continue
148
149            try:
150                # Copy original value if it exists.
151                code_args.append(getattr(code, "co_" + t.cast(str, attr)))
152            except AttributeError:
153                # Some arguments were added later.
154                continue
155
156        code = CodeType(*code_args)
157    except Exception:
158        # Some environments such as Google App Engine don't support
159        # modifying code objects.
160        pass
161
162    # Execute the new code, which is guaranteed to raise, and return
163    # the new traceback without this frame.
164    try:
165        exec(code, globals, locals)
166    except BaseException:
167        return sys.exc_info()[2].tb_next  # type: ignore
168
169
170def get_template_locals(real_locals: t.Mapping[str, t.Any]) -> t.Dict[str, t.Any]:
171    """Based on the runtime locals, get the context that would be
172    available at that point in the template.
173    """
174    # Start with the current template context.
175    ctx: "t.Optional[Context]" = real_locals.get("context")
176
177    if ctx is not None:
178        data: t.Dict[str, t.Any] = ctx.get_all().copy()
179    else:
180        data = {}
181
182    # Might be in a derived context that only sets local variables
183    # rather than pushing a context. Local variables follow the scheme
184    # l_depth_name. Find the highest-depth local that has a value for
185    # each name.
186    local_overrides: t.Dict[str, t.Tuple[int, t.Any]] = {}
187
188    for name, value in real_locals.items():
189        if not name.startswith("l_") or value is missing:
190            # Not a template variable, or no longer relevant.
191            continue
192
193        try:
194            _, depth_str, name = name.split("_", 2)
195            depth = int(depth_str)
196        except ValueError:
197            continue
198
199        cur_depth = local_overrides.get(name, (-1,))[0]
200
201        if cur_depth < depth:
202            local_overrides[name] = (depth, value)
203
204    # Modify the context with any derived context.
205    for name, (_, value) in local_overrides.items():
206        if value is missing:
207            data.pop(name, None)
208        else:
209            data[name] = value
210
211    return data
212
213
214if sys.version_info >= (3, 7):
215    # tb_next is directly assignable as of Python 3.7
216    def tb_set_next(
217        tb: TracebackType, tb_next: t.Optional[TracebackType]
218    ) -> TracebackType:
219        tb.tb_next = tb_next
220        return tb
221
222
223elif platform.python_implementation() == "PyPy":
224    # PyPy might have special support, and won't work with ctypes.
225    try:
226        import tputil  # type: ignore
227    except ImportError:
228        # Without tproxy support, use the original traceback.
229        def tb_set_next(
230            tb: TracebackType, tb_next: t.Optional[TracebackType]
231        ) -> TracebackType:
232            return tb
233
234    else:
235        # With tproxy support, create a proxy around the traceback that
236        # returns the new tb_next.
237        def tb_set_next(
238            tb: TracebackType, tb_next: t.Optional[TracebackType]
239        ) -> TracebackType:
240            def controller(op):  # type: ignore
241                if op.opname == "__getattribute__" and op.args[0] == "tb_next":
242                    return tb_next
243
244                return op.delegate()
245
246            return tputil.make_proxy(controller, obj=tb)  # type: ignore
247
248
249else:
250    # Use ctypes to assign tb_next at the C level since it's read-only
251    # from Python.
252    import ctypes
253
254    class _CTraceback(ctypes.Structure):
255        _fields_ = [
256            # Extra PyObject slots when compiled with Py_TRACE_REFS.
257            ("PyObject_HEAD", ctypes.c_byte * object().__sizeof__()),
258            # Only care about tb_next as an object, not a traceback.
259            ("tb_next", ctypes.py_object),
260        ]
261
262    def tb_set_next(
263        tb: TracebackType, tb_next: t.Optional[TracebackType]
264    ) -> TracebackType:
265        c_tb = _CTraceback.from_address(id(tb))
266
267        # Clear out the old tb_next.
268        if tb.tb_next is not None:
269            c_tb_next = ctypes.py_object(tb.tb_next)
270            c_tb.tb_next = ctypes.py_object()
271            ctypes.pythonapi.Py_DecRef(c_tb_next)
272
273        # Assign the new tb_next.
274        if tb_next is not None:
275            c_tb_next = ctypes.py_object(tb_next)
276            ctypes.pythonapi.Py_IncRef(c_tb_next)
277            c_tb.tb_next = c_tb_next
278
279        return tb
280