formatted_traceback.py 7.69 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
import os
import sys
import importlib
import traceback
from datetime import datetime
from textwrap import indent

from prompt_toolkit.formatted_text import PygmentsTokens
from prompt_toolkit import print_formatted_text

from pygments import lex
from pygments.lexer import RegexLexer, bygroups, default
from pygments.token import Token


class BlissTracebackLexer(RegexLexer):
    """Inspired by pygments.PythonTracebackLexer to colorize tracebacks, but for BlissTraceback
    as well (it handles locals and timestamp)"""

    name = "BlissTraceback"
    tokens = {
        "root": [
            (r"\n", Token.Text),
            (r"\*[^\n]+\n", Token.Comment.Preproc),
            (r"^[0-9]+\/[0-9]+\/[0-9]+ [0-9]+:[0-9]+:[0-9]+\n", Token.Name.Property),
            (r"^Traceback \(most recent call last\):\n", Token.Text, "intb"),
            (
                r"^During handling of the above exception, another exception occurred:\n\n",
                Token.Heading,
            ),
            (
                r"^The above exception was the direct cause of the following exception:\n\n",
                Token.Heading,
            ),
            (r'^(?=  File "[^"]+", line \d+)', Token.Generic.Traceback, "intb"),
            (r"^.*\n", Token.Other),
        ],
        "intb": [
            (
                r'^(  File )("[^"]+")(, line )(\d+)(, in )(.+)(\n)',
                bygroups(
                    Token.Text,
                    Token.Name.Builtin,
                    Token.Text,
                    Token.Number,
                    Token.Text,
                    Token.Name,
                    Token.Text,
                ),
            ),
            (
                r'^(  File )("[^"]+")(, line )(\d+)(\n)',
                bygroups(
                    Token.Text, Token.Name.Builtin, Token.Text, Token.Number, Token.Text
                ),
            ),
            (r"^(?=    @.+\n)", Token.Generic.Traceback, "markers"),
            (
                r"^(    )(.+)(\n)",
                bygroups(Token.Text, Token.Other, Token.Text),
                "markers",
            ),
            (
                r"^([^:]+)(: )(.+)(\n)",
                bygroups(
                    Token.Generic.Error, Token.Text, Token.Name.Exception, Token.Text
                ),
                "#pop",
            ),
            (
                r"^([a-zA-Z_][\w.]*)(:?\n)",
                bygroups(Token.Generic.Error, Token.Text),
                "#pop",
            ),
        ],
        "markers": [
            (r"^(    )\.\.\. \(truncated\)\n", Token.Comment.Preproc),
            (
                r"^(    )(@[^:]+)(: )(\.\.\. \(truncated\)\n)",
                bygroups(
                    Token.Text,
                    Token.Name.Variable,
                    Token.Punctuation,
                    Token.Comment.Preproc,
                ),
            ),
            (
                r"^(    )(@[^:]+)(:)(.+)(\n)",
                bygroups(
                    Token.Text,
                    Token.Name.Variable,
                    Token.Punctuation,
                    Token.Literal.String.Other,
                    Token.Text,
                ),
            ),
            (r"^[\w]*\n", Token.Text, "#pop"),
            default("#pop"),
        ],
    }


def pprint_traceback(formatted_traceback, style):
    """Print a formatted traceback (generic Python traceback or BlissTraceback) with colors,
    using BlissTracebackLexer.
    """
    tokens = list(lex(formatted_traceback, lexer=BlissTracebackLexer()))
    print_formatted_text(PygmentsTokens(tokens), end="", style=style)


class BlissTraceback:
    """Extract traceback content for later formatting without keeping any reference on
    objects to avoid memory leaks. Then the format method can be used to produce various
    formatting of the same traceback.
    """

    _blacklist = [
        # gevent & aiogevent module paths
        os.path.dirname(importlib.util.find_spec("gevent").origin),
        os.path.dirname(importlib.util.find_spec("aiogevent").origin),
        # gevent compiled functions root path
        "src/gevent",
    ]

    def __init__(self, exc_type, exc_value, tb):
        self._datetime = datetime.now()

        # convert traceback to StackSummary to stringify references and avoid memory leak
        self._locals_capture_exc = ""
        try:
            traceback_exc = traceback.TracebackException(
                exc_type, exc_value, tb, capture_locals=True
            )
        except Exception as e:
            # Capture_locals option fails as soon as one local's __repr__ fails.
            # This will be fixed in python 3.11 with the addition of format_locals option,
            # see https://github.com/python/cpython/pull/29299
            # For the moment we can only disable capture_locals
            self._locals_capture_exc = (
                traceback.format_tb(sys.exc_info()[2], -1)[0] + f"  {e}"
            )
            traceback_exc = traceback.TracebackException(
                exc_type, exc_value, tb, capture_locals=False
            )

        self._exc_info = []
        iter = traceback_exc
        while iter is not None:
            exc_type = iter.exc_type
            exc_value = iter._str
            stack = iter.stack
            msg = ""
            if iter.__cause__ is not None:
                msg = "\nThe above exception was the direct cause of the following exception:\n\n"
                iter = iter.__cause__
            elif iter.__context__ is not None:
                msg = "\nDuring handling of the above exception, another exception occurred:\n\n"
                iter = iter.__context__
            else:
                iter = None
            self._exc_info.insert(0, (exc_type, exc_value, stack, msg))

    def _is_file_blacklisted(self, filename):
        for black_path in BlissTraceback._blacklist:
            if filename.startswith(black_path):
                return True
        return False

    def _format_stack(
        self,
        exc_type,
        exc_value,
        stack,
        msg,
        disable_blacklist,
        max_nb_locals,
        max_local_len,
    ):
        text = ""
        for frame in stack:
            if not disable_blacklist and self._is_file_blacklisted(frame.filename):
                continue

            # skip bottom calls from ptpython
            if frame.filename == "<stdin>":
                text = ""

            text += f'  File "{frame.filename}", line {frame.lineno}, in {frame.name}\n'
            if frame._line:
                text += f"    {frame._line}\n"

            if frame.locals is not None:
                for i, (key, val) in enumerate(sorted(frame.locals.items())):
                    if not (max_nb_locals < 0) and i + 1 > max_nb_locals:
                        text += "    ... (truncated)\n"
                        break
                    if len(val) <= max_local_len or max_local_len < 0:
                        text += f"    @{key}: {val}\n"
                    else:
                        text += f"    @{key}: ... (truncated)\n"

        text += f"{exc_type.__name__}: {exc_value}\n"
        return msg + "Traceback (most recent call last):\n" + text

    def format(self, disable_blacklist=False, max_nb_locals=-1, max_local_len=-1):
        timestamp = self._datetime.strftime("%d/%m/%Y %H:%M:%S")
        text = timestamp + "\n"

        if self._locals_capture_exc:
            msg = "Can't display local variables along stack trace, "
            msg += "error occured during recovery:\n"
            msg += self._locals_capture_exc + "\n\n"
            text += indent(msg, "* ")

        # Stack traces formatting
        for exc in self._exc_info:
            text += self._format_stack(
                *exc, disable_blacklist, max_nb_locals, max_local_len
            )
        return text