import os import traceback as traceback_ from types import ModuleType, TracebackType from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Mapping, Optional, Tuple, Type, Union from rich import pretty from rich.columns import Columns from rich.console import RenderResult, group from rich.highlighter import ReprHighlighter from rich.panel import Panel from rich.pretty import Pretty from rich.syntax import PygmentsSyntaxTheme, Syntax from rich.table import Table from rich.text import Text, TextType from rich.traceback import Frame, PathHighlighter, Stack, Trace, LOCALS_MAX_LENGTH, LOCALS_MAX_STRING from rich.traceback import Traceback as BaseTraceback from utils.log._style import MonokaiProStyle if TYPE_CHECKING: from rich.console import ConsoleRenderable # pylint: disable=W0611 __all__ = ("render_scope", "Traceback") def render_scope( scope: Mapping[str, Any], *, title: Optional[TextType] = None, sort_keys: bool = False, indent_guides: bool = False, max_length: Optional[int] = None, max_string: Optional[int] = None, max_depth: Optional[int] = None, ) -> "ConsoleRenderable": """在给定范围内渲染 python 变量 Args: scope (Mapping): 包含变量名称和值的映射. title (str, optional): 标题. 默认为 None. sort_keys (bool, optional): 启用排序. 默认为 True. indent_guides (bool, optional): 启用缩进线. 默认为 False. max_length (int, optional): 缩写前容器的最大长度; 若为 None , 则表示没有缩写. 默认为 None. max_string (int, optional): 截断前字符串的最大长度; 若为 None , 则表示不会截断. 默认为 None. max_depth (int, optional): 嵌套数据结构的最大深度; 若为 None , 则表示会一直递归访问至最后一层. 默认为 None. Returns: ConsoleRenderable: 可被 rich 渲染的对象. """ highlighter = ReprHighlighter() items_table = Table.grid(padding=(0, 1), expand=False) items_table.add_column(justify="right") def sort_items(item: Tuple[str, Any]) -> Tuple[bool, str]: # noinspection PyShadowingNames key, _ = item return not key.startswith("__"), key.lower() # noinspection PyTypeChecker items = sorted(scope.items(), key=sort_items) if sort_keys else scope.items() for key, value in items: key_text = Text.assemble( (key, "scope.key.special" if key.startswith("__") else "scope.key"), (" =", "scope.equals"), ) items_table.add_row( key_text, Pretty( value, highlighter=highlighter, indent_guides=indent_guides, max_length=max_length, max_string=max_string, max_depth=max_depth, ), ) return Panel.fit( items_table, title=title, border_style="scope.border", padding=(0, 1), ) class Traceback(BaseTraceback): locals_max_depth: Optional[int] def __init__(self, *args, locals_max_depth: Optional[int] = None, **kwargs): kwargs.update({"show_locals": True}) super(Traceback, self).__init__(*args, **kwargs) self.locals_max_depth = locals_max_depth @classmethod def from_exception( cls, exc_type: Type[BaseException], exc_value: BaseException, traceback: Optional[TracebackType], *, width: Optional[int] = 100, extra_lines: int = 3, theme: Optional[str] = None, word_wrap: bool = False, show_locals: bool = False, locals_max_length: int = LOCALS_MAX_LENGTH, locals_max_string: int = LOCALS_MAX_STRING, locals_max_depth: Optional[int] = None, locals_hide_dunder: bool = True, locals_hide_sunder: bool = False, indent_guides: bool = True, suppress: Iterable[Union[str, ModuleType]] = (), max_frames: int = 100, ) -> "Traceback": rich_traceback = cls.extract( exc_type=exc_type, exc_value=exc_value, traceback=traceback, show_locals=show_locals, locals_max_depth=locals_max_depth, locals_max_string=locals_max_string, locals_max_length=locals_max_length, ) return cls( rich_traceback, width=width, extra_lines=extra_lines, theme=PygmentsSyntaxTheme(MonokaiProStyle), word_wrap=word_wrap, show_locals=show_locals, indent_guides=indent_guides, locals_max_length=locals_max_length, locals_max_string=locals_max_string, locals_max_depth=locals_max_depth, suppress=suppress, max_frames=max_frames, ) @classmethod def extract( cls, exc_type: Type[BaseException], exc_value: BaseException, traceback: Optional[TracebackType], *, show_locals: bool = False, locals_max_length: int = LOCALS_MAX_LENGTH, locals_max_string: int = LOCALS_MAX_STRING, locals_max_depth: Optional[int] = None, locals_hide_dunder: bool = True, locals_hide_sunder: bool = False, ) -> Trace: # noinspection PyProtectedMember from rich import _IMPORT_CWD stacks: List[Stack] = [] is_cause = False def safe_str(_object: Any) -> str: # noinspection PyBroadException try: return str(_object) except Exception: # pylint: disable=W0703 return "" while True: stack = Stack( exc_type=safe_str(exc_type.__name__), exc_value=safe_str(exc_value), is_cause=is_cause, ) if isinstance(exc_value, SyntaxError): # noinspection PyProtectedMember from rich.traceback import _SyntaxError stack.syntax_error = _SyntaxError( offset=exc_value.offset or 0, filename=exc_value.filename or "?", lineno=exc_value.lineno or 0, line=exc_value.text or "", msg=exc_value.msg, ) stacks.append(stack) append = stack.frames.append for frame_summary, line_no in traceback_.walk_tb(traceback): filename = frame_summary.f_code.co_filename if filename and not filename.startswith("<") and not os.path.isabs(filename): filename = os.path.join(_IMPORT_CWD, filename) if frame_summary.f_locals.get("_rich_traceback_omit", False): continue frame = Frame( filename=filename or "?", lineno=line_no, name=frame_summary.f_code.co_name, locals={ key: pretty.traverse( value, max_length=locals_max_length, max_string=locals_max_string, max_depth=locals_max_depth, ) for key, value in frame_summary.f_locals.items() } if show_locals else None, ) append(frame) if frame_summary.f_locals.get("_rich_traceback_guard", False): del stack.frames[:] cause = getattr(exc_value, "__cause__", None) if cause: exc_type = cause.__class__ exc_value = cause # __traceback__ can be None, e.g. for exceptions raised by the # 'multiprocessing' module traceback = cause.__traceback__ is_cause = True continue cause = exc_value.__context__ if cause and not getattr(exc_value, "__suppress_context__", False): exc_type = cause.__class__ exc_value = cause traceback = cause.__traceback__ is_cause = False continue break trace = Trace(stacks=stacks) return trace @group() def _render_stack(self, stack: Stack) -> RenderResult: path_highlighter = PathHighlighter() theme = self.theme code_cache: Dict[str, str] = {} # noinspection PyShadowingNames def read_code(filename: str) -> str: code = code_cache.get(filename) if code is None: with open(filename, "rt", encoding="utf-8", errors="replace") as code_file: code = code_file.read() code_cache[filename] = code return code # noinspection PyShadowingNames def render_locals(frame: Frame) -> Iterable["ConsoleRenderable"]: if frame.locals: yield render_scope( scope=frame.locals, title="locals", indent_guides=self.indent_guides, max_length=self.locals_max_length, max_string=self.locals_max_string, max_depth=self.locals_max_depth, ) exclude_frames: Optional[range] = None if self.max_frames != 0: exclude_frames = range( self.max_frames // 2, len(stack.frames) - self.max_frames // 2, ) excluded = False for frame_index, frame in enumerate(stack.frames): if exclude_frames and frame_index in exclude_frames: excluded = True continue if excluded: if exclude_frames is None: raise ValueError(exclude_frames) yield Text( f"\n... {len(exclude_frames)} frames hidden ...", justify="center", style="traceback.error", ) excluded = False first = frame_index == 0 frame_filename = frame.filename suppressed = any(frame_filename.startswith(path) for path in self.suppress) text = Text.assemble( path_highlighter(Text(frame.filename, style="pygments.string")), (":", "pygments.text"), (str(frame.lineno), "pygments.number"), " in ", (frame.name, "pygments.function"), style="pygments.text", ) if not frame.filename.startswith("<") and not first: yield "" yield text if frame.filename.startswith("<"): yield from render_locals(frame) continue if not suppressed: try: if self.width is not None: code_width = self.width - 5 else: code_width = 100 code = read_code(frame.filename) lexer_name = self._guess_lexer(frame.filename, code) syntax = Syntax( code, lexer_name, theme=theme, line_numbers=True, line_range=( frame.lineno - self.extra_lines, frame.lineno + self.extra_lines, ), highlight_lines={frame.lineno}, word_wrap=self.word_wrap, code_width=code_width, indent_guides=self.indent_guides, dedent=False, ) yield "" except Exception as error: # pylint: disable=W0703 yield Text.assemble( (f"\n{error}", "traceback.error"), ) else: yield ( Columns( [ syntax, *render_locals(frame), ], padding=1, ) if frame.locals else syntax )