import os import traceback as traceback_ from types import ( ModuleType, TracebackType, ) from typing import ( Any, Dict, Iterable, List, Mapping, Optional, TYPE_CHECKING, 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, 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 = True, indent_guides: bool = True, locals_max_length: int = 10, locals_max_string: int = 80, locals_max_depth: Optional[int] = None, 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 = 10, locals_max_string: int = 80, locals_max_depth: Optional[int] = None, ) -> 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 )