996 lines
35 KiB
Python
996 lines
35 KiB
Python
|
import builtins
|
||
|
import collections
|
||
|
import dataclasses
|
||
|
import inspect
|
||
|
import os
|
||
|
import sys
|
||
|
from array import array
|
||
|
from collections import Counter, UserDict, UserList, defaultdict, deque
|
||
|
from dataclasses import dataclass, fields, is_dataclass
|
||
|
from inspect import isclass
|
||
|
from itertools import islice
|
||
|
from types import MappingProxyType
|
||
|
from typing import (
|
||
|
TYPE_CHECKING,
|
||
|
Any,
|
||
|
Callable,
|
||
|
DefaultDict,
|
||
|
Dict,
|
||
|
Iterable,
|
||
|
List,
|
||
|
Optional,
|
||
|
Sequence,
|
||
|
Set,
|
||
|
Tuple,
|
||
|
Union,
|
||
|
)
|
||
|
|
||
|
from rich.repr import RichReprResult
|
||
|
|
||
|
try:
|
||
|
import attr as _attr_module
|
||
|
|
||
|
_has_attrs = hasattr(_attr_module, "ib")
|
||
|
except ImportError: # pragma: no cover
|
||
|
_has_attrs = False
|
||
|
|
||
|
from . import get_console
|
||
|
from ._loop import loop_last
|
||
|
from ._pick import pick_bool
|
||
|
from .abc import RichRenderable
|
||
|
from .cells import cell_len
|
||
|
from .highlighter import ReprHighlighter
|
||
|
from .jupyter import JupyterMixin, JupyterRenderable
|
||
|
from .measure import Measurement
|
||
|
from .text import Text
|
||
|
|
||
|
if TYPE_CHECKING:
|
||
|
from .console import (
|
||
|
Console,
|
||
|
ConsoleOptions,
|
||
|
HighlighterType,
|
||
|
JustifyMethod,
|
||
|
OverflowMethod,
|
||
|
RenderResult,
|
||
|
)
|
||
|
|
||
|
|
||
|
def _is_attr_object(obj: Any) -> bool:
|
||
|
"""Check if an object was created with attrs module."""
|
||
|
return _has_attrs and _attr_module.has(type(obj))
|
||
|
|
||
|
|
||
|
def _get_attr_fields(obj: Any) -> Sequence["_attr_module.Attribute[Any]"]:
|
||
|
"""Get fields for an attrs object."""
|
||
|
return _attr_module.fields(type(obj)) if _has_attrs else []
|
||
|
|
||
|
|
||
|
def _is_dataclass_repr(obj: object) -> bool:
|
||
|
"""Check if an instance of a dataclass contains the default repr.
|
||
|
|
||
|
Args:
|
||
|
obj (object): A dataclass instance.
|
||
|
|
||
|
Returns:
|
||
|
bool: True if the default repr is used, False if there is a custom repr.
|
||
|
"""
|
||
|
# Digging in to a lot of internals here
|
||
|
# Catching all exceptions in case something is missing on a non CPython implementation
|
||
|
try:
|
||
|
return obj.__repr__.__code__.co_filename == dataclasses.__file__
|
||
|
except Exception: # pragma: no coverage
|
||
|
return False
|
||
|
|
||
|
|
||
|
_dummy_namedtuple = collections.namedtuple("_dummy_namedtuple", [])
|
||
|
|
||
|
|
||
|
def _has_default_namedtuple_repr(obj: object) -> bool:
|
||
|
"""Check if an instance of namedtuple contains the default repr
|
||
|
|
||
|
Args:
|
||
|
obj (object): A namedtuple
|
||
|
|
||
|
Returns:
|
||
|
bool: True if the default repr is used, False if there's a custom repr.
|
||
|
"""
|
||
|
obj_file = None
|
||
|
try:
|
||
|
obj_file = inspect.getfile(obj.__repr__)
|
||
|
except (OSError, TypeError):
|
||
|
# OSError handles case where object is defined in __main__ scope, e.g. REPL - no filename available.
|
||
|
# TypeError trapped defensively, in case of object without filename slips through.
|
||
|
pass
|
||
|
default_repr_file = inspect.getfile(_dummy_namedtuple.__repr__)
|
||
|
return obj_file == default_repr_file
|
||
|
|
||
|
|
||
|
def _ipy_display_hook(
|
||
|
value: Any,
|
||
|
console: Optional["Console"] = None,
|
||
|
overflow: "OverflowMethod" = "ignore",
|
||
|
crop: bool = False,
|
||
|
indent_guides: bool = False,
|
||
|
max_length: Optional[int] = None,
|
||
|
max_string: Optional[int] = None,
|
||
|
max_depth: Optional[int] = None,
|
||
|
expand_all: bool = False,
|
||
|
) -> Union[str, None]:
|
||
|
# needed here to prevent circular import:
|
||
|
from .console import ConsoleRenderable
|
||
|
|
||
|
# always skip rich generated jupyter renderables or None values
|
||
|
if _safe_isinstance(value, JupyterRenderable) or value is None:
|
||
|
return None
|
||
|
|
||
|
console = console or get_console()
|
||
|
|
||
|
with console.capture() as capture:
|
||
|
# certain renderables should start on a new line
|
||
|
if _safe_isinstance(value, ConsoleRenderable):
|
||
|
console.line()
|
||
|
console.print(
|
||
|
value
|
||
|
if _safe_isinstance(value, RichRenderable)
|
||
|
else Pretty(
|
||
|
value,
|
||
|
overflow=overflow,
|
||
|
indent_guides=indent_guides,
|
||
|
max_length=max_length,
|
||
|
max_string=max_string,
|
||
|
max_depth=max_depth,
|
||
|
expand_all=expand_all,
|
||
|
margin=12,
|
||
|
),
|
||
|
crop=crop,
|
||
|
new_line_start=True,
|
||
|
end="",
|
||
|
)
|
||
|
# strip trailing newline, not usually part of a text repr
|
||
|
# I'm not sure if this should be prevented at a lower level
|
||
|
return capture.get().rstrip("\n")
|
||
|
|
||
|
|
||
|
def _safe_isinstance(
|
||
|
obj: object, class_or_tuple: Union[type, Tuple[type, ...]]
|
||
|
) -> bool:
|
||
|
"""isinstance can fail in rare cases, for example types with no __class__"""
|
||
|
try:
|
||
|
return isinstance(obj, class_or_tuple)
|
||
|
except Exception:
|
||
|
return False
|
||
|
|
||
|
|
||
|
def install(
|
||
|
console: Optional["Console"] = None,
|
||
|
overflow: "OverflowMethod" = "ignore",
|
||
|
crop: bool = False,
|
||
|
indent_guides: bool = False,
|
||
|
max_length: Optional[int] = None,
|
||
|
max_string: Optional[int] = None,
|
||
|
max_depth: Optional[int] = None,
|
||
|
expand_all: bool = False,
|
||
|
) -> None:
|
||
|
"""Install automatic pretty printing in the Python REPL.
|
||
|
|
||
|
Args:
|
||
|
console (Console, optional): Console instance or ``None`` to use global console. Defaults to None.
|
||
|
overflow (Optional[OverflowMethod], optional): Overflow method. Defaults to "ignore".
|
||
|
crop (Optional[bool], optional): Enable cropping of long lines. Defaults to False.
|
||
|
indent_guides (bool, optional): Enable indentation guides. Defaults to False.
|
||
|
max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
|
||
|
Defaults to None.
|
||
|
max_string (int, optional): Maximum length of string before truncating, or None to disable. Defaults to None.
|
||
|
max_depth (int, optional): Maximum depth of nested data structures, or None for no maximum. Defaults to None.
|
||
|
expand_all (bool, optional): Expand all containers. Defaults to False.
|
||
|
max_frames (int): Maximum number of frames to show in a traceback, 0 for no maximum. Defaults to 100.
|
||
|
"""
|
||
|
from rich import get_console
|
||
|
|
||
|
console = console or get_console()
|
||
|
assert console is not None
|
||
|
|
||
|
def display_hook(value: Any) -> None:
|
||
|
"""Replacement sys.displayhook which prettifies objects with Rich."""
|
||
|
if value is not None:
|
||
|
assert console is not None
|
||
|
builtins._ = None # type: ignore[attr-defined]
|
||
|
console.print(
|
||
|
value
|
||
|
if _safe_isinstance(value, RichRenderable)
|
||
|
else Pretty(
|
||
|
value,
|
||
|
overflow=overflow,
|
||
|
indent_guides=indent_guides,
|
||
|
max_length=max_length,
|
||
|
max_string=max_string,
|
||
|
max_depth=max_depth,
|
||
|
expand_all=expand_all,
|
||
|
),
|
||
|
crop=crop,
|
||
|
)
|
||
|
builtins._ = value # type: ignore[attr-defined]
|
||
|
|
||
|
try:
|
||
|
ip = get_ipython() # type: ignore[name-defined]
|
||
|
except NameError:
|
||
|
sys.displayhook = display_hook
|
||
|
else:
|
||
|
from IPython.core.formatters import BaseFormatter
|
||
|
|
||
|
class RichFormatter(BaseFormatter): # type: ignore[misc]
|
||
|
pprint: bool = True
|
||
|
|
||
|
def __call__(self, value: Any) -> Any:
|
||
|
if self.pprint:
|
||
|
return _ipy_display_hook(
|
||
|
value,
|
||
|
console=get_console(),
|
||
|
overflow=overflow,
|
||
|
indent_guides=indent_guides,
|
||
|
max_length=max_length,
|
||
|
max_string=max_string,
|
||
|
max_depth=max_depth,
|
||
|
expand_all=expand_all,
|
||
|
)
|
||
|
else:
|
||
|
return repr(value)
|
||
|
|
||
|
# replace plain text formatter with rich formatter
|
||
|
rich_formatter = RichFormatter()
|
||
|
ip.display_formatter.formatters["text/plain"] = rich_formatter
|
||
|
|
||
|
|
||
|
class Pretty(JupyterMixin):
|
||
|
"""A rich renderable that pretty prints an object.
|
||
|
|
||
|
Args:
|
||
|
_object (Any): An object to pretty print.
|
||
|
highlighter (HighlighterType, optional): Highlighter object to apply to result, or None for ReprHighlighter. Defaults to None.
|
||
|
indent_size (int, optional): Number of spaces in indent. Defaults to 4.
|
||
|
justify (JustifyMethod, optional): Justify method, or None for default. Defaults to None.
|
||
|
overflow (OverflowMethod, optional): Overflow method, or None for default. Defaults to None.
|
||
|
no_wrap (Optional[bool], optional): Disable word wrapping. Defaults to False.
|
||
|
indent_guides (bool, optional): Enable indentation guides. Defaults to False.
|
||
|
max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
|
||
|
Defaults to None.
|
||
|
max_string (int, optional): Maximum length of string before truncating, or None to disable. Defaults to None.
|
||
|
max_depth (int, optional): Maximum depth of nested data structures, or None for no maximum. Defaults to None.
|
||
|
expand_all (bool, optional): Expand all containers. Defaults to False.
|
||
|
margin (int, optional): Subtrace a margin from width to force containers to expand earlier. Defaults to 0.
|
||
|
insert_line (bool, optional): Insert a new line if the output has multiple new lines. Defaults to False.
|
||
|
"""
|
||
|
|
||
|
def __init__(
|
||
|
self,
|
||
|
_object: Any,
|
||
|
highlighter: Optional["HighlighterType"] = None,
|
||
|
*,
|
||
|
indent_size: int = 4,
|
||
|
justify: Optional["JustifyMethod"] = None,
|
||
|
overflow: Optional["OverflowMethod"] = None,
|
||
|
no_wrap: Optional[bool] = False,
|
||
|
indent_guides: bool = False,
|
||
|
max_length: Optional[int] = None,
|
||
|
max_string: Optional[int] = None,
|
||
|
max_depth: Optional[int] = None,
|
||
|
expand_all: bool = False,
|
||
|
margin: int = 0,
|
||
|
insert_line: bool = False,
|
||
|
) -> None:
|
||
|
self._object = _object
|
||
|
self.highlighter = highlighter or ReprHighlighter()
|
||
|
self.indent_size = indent_size
|
||
|
self.justify: Optional["JustifyMethod"] = justify
|
||
|
self.overflow: Optional["OverflowMethod"] = overflow
|
||
|
self.no_wrap = no_wrap
|
||
|
self.indent_guides = indent_guides
|
||
|
self.max_length = max_length
|
||
|
self.max_string = max_string
|
||
|
self.max_depth = max_depth
|
||
|
self.expand_all = expand_all
|
||
|
self.margin = margin
|
||
|
self.insert_line = insert_line
|
||
|
|
||
|
def __rich_console__(
|
||
|
self, console: "Console", options: "ConsoleOptions"
|
||
|
) -> "RenderResult":
|
||
|
pretty_str = pretty_repr(
|
||
|
self._object,
|
||
|
max_width=options.max_width - self.margin,
|
||
|
indent_size=self.indent_size,
|
||
|
max_length=self.max_length,
|
||
|
max_string=self.max_string,
|
||
|
max_depth=self.max_depth,
|
||
|
expand_all=self.expand_all,
|
||
|
)
|
||
|
pretty_text = Text.from_ansi(
|
||
|
pretty_str,
|
||
|
justify=self.justify or options.justify,
|
||
|
overflow=self.overflow or options.overflow,
|
||
|
no_wrap=pick_bool(self.no_wrap, options.no_wrap),
|
||
|
style="pretty",
|
||
|
)
|
||
|
pretty_text = (
|
||
|
self.highlighter(pretty_text)
|
||
|
if pretty_text
|
||
|
else Text(
|
||
|
f"{type(self._object)}.__repr__ returned empty string",
|
||
|
style="dim italic",
|
||
|
)
|
||
|
)
|
||
|
if self.indent_guides and not options.ascii_only:
|
||
|
pretty_text = pretty_text.with_indent_guides(
|
||
|
self.indent_size, style="repr.indent"
|
||
|
)
|
||
|
if self.insert_line and "\n" in pretty_text:
|
||
|
yield ""
|
||
|
yield pretty_text
|
||
|
|
||
|
def __rich_measure__(
|
||
|
self, console: "Console", options: "ConsoleOptions"
|
||
|
) -> "Measurement":
|
||
|
pretty_str = pretty_repr(
|
||
|
self._object,
|
||
|
max_width=options.max_width,
|
||
|
indent_size=self.indent_size,
|
||
|
max_length=self.max_length,
|
||
|
max_string=self.max_string,
|
||
|
max_depth=self.max_depth,
|
||
|
expand_all=self.expand_all,
|
||
|
)
|
||
|
text_width = (
|
||
|
max(cell_len(line) for line in pretty_str.splitlines()) if pretty_str else 0
|
||
|
)
|
||
|
return Measurement(text_width, text_width)
|
||
|
|
||
|
|
||
|
def _get_braces_for_defaultdict(_object: DefaultDict[Any, Any]) -> Tuple[str, str, str]:
|
||
|
return (
|
||
|
f"defaultdict({_object.default_factory!r}, {{",
|
||
|
"})",
|
||
|
f"defaultdict({_object.default_factory!r}, {{}})",
|
||
|
)
|
||
|
|
||
|
|
||
|
def _get_braces_for_array(_object: "array[Any]") -> Tuple[str, str, str]:
|
||
|
return (f"array({_object.typecode!r}, [", "])", f"array({_object.typecode!r})")
|
||
|
|
||
|
|
||
|
_BRACES: Dict[type, Callable[[Any], Tuple[str, str, str]]] = {
|
||
|
os._Environ: lambda _object: ("environ({", "})", "environ({})"),
|
||
|
array: _get_braces_for_array,
|
||
|
defaultdict: _get_braces_for_defaultdict,
|
||
|
Counter: lambda _object: ("Counter({", "})", "Counter()"),
|
||
|
deque: lambda _object: ("deque([", "])", "deque()"),
|
||
|
dict: lambda _object: ("{", "}", "{}"),
|
||
|
UserDict: lambda _object: ("{", "}", "{}"),
|
||
|
frozenset: lambda _object: ("frozenset({", "})", "frozenset()"),
|
||
|
list: lambda _object: ("[", "]", "[]"),
|
||
|
UserList: lambda _object: ("[", "]", "[]"),
|
||
|
set: lambda _object: ("{", "}", "set()"),
|
||
|
tuple: lambda _object: ("(", ")", "()"),
|
||
|
MappingProxyType: lambda _object: ("mappingproxy({", "})", "mappingproxy({})"),
|
||
|
}
|
||
|
_CONTAINERS = tuple(_BRACES.keys())
|
||
|
_MAPPING_CONTAINERS = (dict, os._Environ, MappingProxyType, UserDict)
|
||
|
|
||
|
|
||
|
def is_expandable(obj: Any) -> bool:
|
||
|
"""Check if an object may be expanded by pretty print."""
|
||
|
return (
|
||
|
_safe_isinstance(obj, _CONTAINERS)
|
||
|
or (is_dataclass(obj))
|
||
|
or (hasattr(obj, "__rich_repr__"))
|
||
|
or _is_attr_object(obj)
|
||
|
) and not isclass(obj)
|
||
|
|
||
|
|
||
|
@dataclass
|
||
|
class Node:
|
||
|
"""A node in a repr tree. May be atomic or a container."""
|
||
|
|
||
|
key_repr: str = ""
|
||
|
value_repr: str = ""
|
||
|
open_brace: str = ""
|
||
|
close_brace: str = ""
|
||
|
empty: str = ""
|
||
|
last: bool = False
|
||
|
is_tuple: bool = False
|
||
|
is_namedtuple: bool = False
|
||
|
children: Optional[List["Node"]] = None
|
||
|
key_separator: str = ": "
|
||
|
separator: str = ", "
|
||
|
|
||
|
def iter_tokens(self) -> Iterable[str]:
|
||
|
"""Generate tokens for this node."""
|
||
|
if self.key_repr:
|
||
|
yield self.key_repr
|
||
|
yield self.key_separator
|
||
|
if self.value_repr:
|
||
|
yield self.value_repr
|
||
|
elif self.children is not None:
|
||
|
if self.children:
|
||
|
yield self.open_brace
|
||
|
if self.is_tuple and not self.is_namedtuple and len(self.children) == 1:
|
||
|
yield from self.children[0].iter_tokens()
|
||
|
yield ","
|
||
|
else:
|
||
|
for child in self.children:
|
||
|
yield from child.iter_tokens()
|
||
|
if not child.last:
|
||
|
yield self.separator
|
||
|
yield self.close_brace
|
||
|
else:
|
||
|
yield self.empty
|
||
|
|
||
|
def check_length(self, start_length: int, max_length: int) -> bool:
|
||
|
"""Check the length fits within a limit.
|
||
|
|
||
|
Args:
|
||
|
start_length (int): Starting length of the line (indent, prefix, suffix).
|
||
|
max_length (int): Maximum length.
|
||
|
|
||
|
Returns:
|
||
|
bool: True if the node can be rendered within max length, otherwise False.
|
||
|
"""
|
||
|
total_length = start_length
|
||
|
for token in self.iter_tokens():
|
||
|
total_length += cell_len(token)
|
||
|
if total_length > max_length:
|
||
|
return False
|
||
|
return True
|
||
|
|
||
|
def __str__(self) -> str:
|
||
|
repr_text = "".join(self.iter_tokens())
|
||
|
return repr_text
|
||
|
|
||
|
def render(
|
||
|
self, max_width: int = 80, indent_size: int = 4, expand_all: bool = False
|
||
|
) -> str:
|
||
|
"""Render the node to a pretty repr.
|
||
|
|
||
|
Args:
|
||
|
max_width (int, optional): Maximum width of the repr. Defaults to 80.
|
||
|
indent_size (int, optional): Size of indents. Defaults to 4.
|
||
|
expand_all (bool, optional): Expand all levels. Defaults to False.
|
||
|
|
||
|
Returns:
|
||
|
str: A repr string of the original object.
|
||
|
"""
|
||
|
lines = [_Line(node=self, is_root=True)]
|
||
|
line_no = 0
|
||
|
while line_no < len(lines):
|
||
|
line = lines[line_no]
|
||
|
if line.expandable and not line.expanded:
|
||
|
if expand_all or not line.check_length(max_width):
|
||
|
lines[line_no : line_no + 1] = line.expand(indent_size)
|
||
|
line_no += 1
|
||
|
|
||
|
repr_str = "\n".join(str(line) for line in lines)
|
||
|
return repr_str
|
||
|
|
||
|
|
||
|
@dataclass
|
||
|
class _Line:
|
||
|
"""A line in repr output."""
|
||
|
|
||
|
parent: Optional["_Line"] = None
|
||
|
is_root: bool = False
|
||
|
node: Optional[Node] = None
|
||
|
text: str = ""
|
||
|
suffix: str = ""
|
||
|
whitespace: str = ""
|
||
|
expanded: bool = False
|
||
|
last: bool = False
|
||
|
|
||
|
@property
|
||
|
def expandable(self) -> bool:
|
||
|
"""Check if the line may be expanded."""
|
||
|
return bool(self.node is not None and self.node.children)
|
||
|
|
||
|
def check_length(self, max_length: int) -> bool:
|
||
|
"""Check this line fits within a given number of cells."""
|
||
|
start_length = (
|
||
|
len(self.whitespace) + cell_len(self.text) + cell_len(self.suffix)
|
||
|
)
|
||
|
assert self.node is not None
|
||
|
return self.node.check_length(start_length, max_length)
|
||
|
|
||
|
def expand(self, indent_size: int) -> Iterable["_Line"]:
|
||
|
"""Expand this line by adding children on their own line."""
|
||
|
node = self.node
|
||
|
assert node is not None
|
||
|
whitespace = self.whitespace
|
||
|
assert node.children
|
||
|
if node.key_repr:
|
||
|
new_line = yield _Line(
|
||
|
text=f"{node.key_repr}{node.key_separator}{node.open_brace}",
|
||
|
whitespace=whitespace,
|
||
|
)
|
||
|
else:
|
||
|
new_line = yield _Line(text=node.open_brace, whitespace=whitespace)
|
||
|
child_whitespace = self.whitespace + " " * indent_size
|
||
|
tuple_of_one = node.is_tuple and len(node.children) == 1
|
||
|
for last, child in loop_last(node.children):
|
||
|
separator = "," if tuple_of_one else node.separator
|
||
|
line = _Line(
|
||
|
parent=new_line,
|
||
|
node=child,
|
||
|
whitespace=child_whitespace,
|
||
|
suffix=separator,
|
||
|
last=last and not tuple_of_one,
|
||
|
)
|
||
|
yield line
|
||
|
|
||
|
yield _Line(
|
||
|
text=node.close_brace,
|
||
|
whitespace=whitespace,
|
||
|
suffix=self.suffix,
|
||
|
last=self.last,
|
||
|
)
|
||
|
|
||
|
def __str__(self) -> str:
|
||
|
if self.last:
|
||
|
return f"{self.whitespace}{self.text}{self.node or ''}"
|
||
|
else:
|
||
|
return (
|
||
|
f"{self.whitespace}{self.text}{self.node or ''}{self.suffix.rstrip()}"
|
||
|
)
|
||
|
|
||
|
|
||
|
def _is_namedtuple(obj: Any) -> bool:
|
||
|
"""Checks if an object is most likely a namedtuple. It is possible
|
||
|
to craft an object that passes this check and isn't a namedtuple, but
|
||
|
there is only a minuscule chance of this happening unintentionally.
|
||
|
|
||
|
Args:
|
||
|
obj (Any): The object to test
|
||
|
|
||
|
Returns:
|
||
|
bool: True if the object is a namedtuple. False otherwise.
|
||
|
"""
|
||
|
try:
|
||
|
fields = getattr(obj, "_fields", None)
|
||
|
except Exception:
|
||
|
# Being very defensive - if we cannot get the attr then its not a namedtuple
|
||
|
return False
|
||
|
return isinstance(obj, tuple) and isinstance(fields, tuple)
|
||
|
|
||
|
|
||
|
def traverse(
|
||
|
_object: Any,
|
||
|
max_length: Optional[int] = None,
|
||
|
max_string: Optional[int] = None,
|
||
|
max_depth: Optional[int] = None,
|
||
|
) -> Node:
|
||
|
"""Traverse object and generate a tree.
|
||
|
|
||
|
Args:
|
||
|
_object (Any): Object to be traversed.
|
||
|
max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
|
||
|
Defaults to None.
|
||
|
max_string (int, optional): Maximum length of string before truncating, or None to disable truncating.
|
||
|
Defaults to None.
|
||
|
max_depth (int, optional): Maximum depth of data structures, or None for no maximum.
|
||
|
Defaults to None.
|
||
|
|
||
|
Returns:
|
||
|
Node: The root of a tree structure which can be used to render a pretty repr.
|
||
|
"""
|
||
|
|
||
|
def to_repr(obj: Any) -> str:
|
||
|
"""Get repr string for an object, but catch errors."""
|
||
|
if (
|
||
|
max_string is not None
|
||
|
and _safe_isinstance(obj, (bytes, str))
|
||
|
and len(obj) > max_string
|
||
|
):
|
||
|
truncated = len(obj) - max_string
|
||
|
obj_repr = f"{obj[:max_string]!r}+{truncated}"
|
||
|
else:
|
||
|
try:
|
||
|
obj_repr = repr(obj)
|
||
|
except Exception as error:
|
||
|
obj_repr = f"<repr-error {str(error)!r}>"
|
||
|
return obj_repr
|
||
|
|
||
|
visited_ids: Set[int] = set()
|
||
|
push_visited = visited_ids.add
|
||
|
pop_visited = visited_ids.remove
|
||
|
|
||
|
def _traverse(obj: Any, root: bool = False, depth: int = 0) -> Node:
|
||
|
"""Walk the object depth first."""
|
||
|
|
||
|
obj_id = id(obj)
|
||
|
if obj_id in visited_ids:
|
||
|
# Recursion detected
|
||
|
return Node(value_repr="...")
|
||
|
|
||
|
obj_type = type(obj)
|
||
|
children: List[Node]
|
||
|
reached_max_depth = max_depth is not None and depth >= max_depth
|
||
|
|
||
|
def iter_rich_args(rich_args: Any) -> Iterable[Union[Any, Tuple[str, Any]]]:
|
||
|
for arg in rich_args:
|
||
|
if _safe_isinstance(arg, tuple):
|
||
|
if len(arg) == 3:
|
||
|
key, child, default = arg
|
||
|
if default == child:
|
||
|
continue
|
||
|
yield key, child
|
||
|
elif len(arg) == 2:
|
||
|
key, child = arg
|
||
|
yield key, child
|
||
|
elif len(arg) == 1:
|
||
|
yield arg[0]
|
||
|
else:
|
||
|
yield arg
|
||
|
|
||
|
try:
|
||
|
fake_attributes = hasattr(
|
||
|
obj, "awehoi234_wdfjwljet234_234wdfoijsdfmmnxpi492"
|
||
|
)
|
||
|
except Exception:
|
||
|
fake_attributes = False
|
||
|
|
||
|
rich_repr_result: Optional[RichReprResult] = None
|
||
|
if not fake_attributes:
|
||
|
try:
|
||
|
if hasattr(obj, "__rich_repr__") and not isclass(obj):
|
||
|
rich_repr_result = obj.__rich_repr__()
|
||
|
except Exception:
|
||
|
pass
|
||
|
|
||
|
if rich_repr_result is not None:
|
||
|
push_visited(obj_id)
|
||
|
angular = getattr(obj.__rich_repr__, "angular", False)
|
||
|
args = list(iter_rich_args(rich_repr_result))
|
||
|
class_name = obj.__class__.__name__
|
||
|
|
||
|
if args:
|
||
|
children = []
|
||
|
append = children.append
|
||
|
|
||
|
if reached_max_depth:
|
||
|
if angular:
|
||
|
node = Node(value_repr=f"<{class_name}...>")
|
||
|
else:
|
||
|
node = Node(value_repr=f"{class_name}(...)")
|
||
|
else:
|
||
|
if angular:
|
||
|
node = Node(
|
||
|
open_brace=f"<{class_name} ",
|
||
|
close_brace=">",
|
||
|
children=children,
|
||
|
last=root,
|
||
|
separator=" ",
|
||
|
)
|
||
|
else:
|
||
|
node = Node(
|
||
|
open_brace=f"{class_name}(",
|
||
|
close_brace=")",
|
||
|
children=children,
|
||
|
last=root,
|
||
|
)
|
||
|
for last, arg in loop_last(args):
|
||
|
if _safe_isinstance(arg, tuple):
|
||
|
key, child = arg
|
||
|
child_node = _traverse(child, depth=depth + 1)
|
||
|
child_node.last = last
|
||
|
child_node.key_repr = key
|
||
|
child_node.key_separator = "="
|
||
|
append(child_node)
|
||
|
else:
|
||
|
child_node = _traverse(arg, depth=depth + 1)
|
||
|
child_node.last = last
|
||
|
append(child_node)
|
||
|
else:
|
||
|
node = Node(
|
||
|
value_repr=f"<{class_name}>" if angular else f"{class_name}()",
|
||
|
children=[],
|
||
|
last=root,
|
||
|
)
|
||
|
pop_visited(obj_id)
|
||
|
elif _is_attr_object(obj) and not fake_attributes:
|
||
|
push_visited(obj_id)
|
||
|
children = []
|
||
|
append = children.append
|
||
|
|
||
|
attr_fields = _get_attr_fields(obj)
|
||
|
if attr_fields:
|
||
|
if reached_max_depth:
|
||
|
node = Node(value_repr=f"{obj.__class__.__name__}(...)")
|
||
|
else:
|
||
|
node = Node(
|
||
|
open_brace=f"{obj.__class__.__name__}(",
|
||
|
close_brace=")",
|
||
|
children=children,
|
||
|
last=root,
|
||
|
)
|
||
|
|
||
|
def iter_attrs() -> (
|
||
|
Iterable[Tuple[str, Any, Optional[Callable[[Any], str]]]]
|
||
|
):
|
||
|
"""Iterate over attr fields and values."""
|
||
|
for attr in attr_fields:
|
||
|
if attr.repr:
|
||
|
try:
|
||
|
value = getattr(obj, attr.name)
|
||
|
except Exception as error:
|
||
|
# Can happen, albeit rarely
|
||
|
yield (attr.name, error, None)
|
||
|
else:
|
||
|
yield (
|
||
|
attr.name,
|
||
|
value,
|
||
|
attr.repr if callable(attr.repr) else None,
|
||
|
)
|
||
|
|
||
|
for last, (name, value, repr_callable) in loop_last(iter_attrs()):
|
||
|
if repr_callable:
|
||
|
child_node = Node(value_repr=str(repr_callable(value)))
|
||
|
else:
|
||
|
child_node = _traverse(value, depth=depth + 1)
|
||
|
child_node.last = last
|
||
|
child_node.key_repr = name
|
||
|
child_node.key_separator = "="
|
||
|
append(child_node)
|
||
|
else:
|
||
|
node = Node(
|
||
|
value_repr=f"{obj.__class__.__name__}()", children=[], last=root
|
||
|
)
|
||
|
pop_visited(obj_id)
|
||
|
elif (
|
||
|
is_dataclass(obj)
|
||
|
and not _safe_isinstance(obj, type)
|
||
|
and not fake_attributes
|
||
|
and _is_dataclass_repr(obj)
|
||
|
):
|
||
|
push_visited(obj_id)
|
||
|
children = []
|
||
|
append = children.append
|
||
|
if reached_max_depth:
|
||
|
node = Node(value_repr=f"{obj.__class__.__name__}(...)")
|
||
|
else:
|
||
|
node = Node(
|
||
|
open_brace=f"{obj.__class__.__name__}(",
|
||
|
close_brace=")",
|
||
|
children=children,
|
||
|
last=root,
|
||
|
empty=f"{obj.__class__.__name__}()",
|
||
|
)
|
||
|
|
||
|
for last, field in loop_last(
|
||
|
field for field in fields(obj) if field.repr
|
||
|
):
|
||
|
child_node = _traverse(getattr(obj, field.name), depth=depth + 1)
|
||
|
child_node.key_repr = field.name
|
||
|
child_node.last = last
|
||
|
child_node.key_separator = "="
|
||
|
append(child_node)
|
||
|
|
||
|
pop_visited(obj_id)
|
||
|
elif _is_namedtuple(obj) and _has_default_namedtuple_repr(obj):
|
||
|
push_visited(obj_id)
|
||
|
class_name = obj.__class__.__name__
|
||
|
if reached_max_depth:
|
||
|
# If we've reached the max depth, we still show the class name, but not its contents
|
||
|
node = Node(
|
||
|
value_repr=f"{class_name}(...)",
|
||
|
)
|
||
|
else:
|
||
|
children = []
|
||
|
append = children.append
|
||
|
node = Node(
|
||
|
open_brace=f"{class_name}(",
|
||
|
close_brace=")",
|
||
|
children=children,
|
||
|
empty=f"{class_name}()",
|
||
|
)
|
||
|
for last, (key, value) in loop_last(obj._asdict().items()):
|
||
|
child_node = _traverse(value, depth=depth + 1)
|
||
|
child_node.key_repr = key
|
||
|
child_node.last = last
|
||
|
child_node.key_separator = "="
|
||
|
append(child_node)
|
||
|
pop_visited(obj_id)
|
||
|
elif _safe_isinstance(obj, _CONTAINERS):
|
||
|
for container_type in _CONTAINERS:
|
||
|
if _safe_isinstance(obj, container_type):
|
||
|
obj_type = container_type
|
||
|
break
|
||
|
|
||
|
push_visited(obj_id)
|
||
|
|
||
|
open_brace, close_brace, empty = _BRACES[obj_type](obj)
|
||
|
|
||
|
if reached_max_depth:
|
||
|
node = Node(value_repr=f"{open_brace}...{close_brace}")
|
||
|
elif obj_type.__repr__ != type(obj).__repr__:
|
||
|
node = Node(value_repr=to_repr(obj), last=root)
|
||
|
elif obj:
|
||
|
children = []
|
||
|
node = Node(
|
||
|
open_brace=open_brace,
|
||
|
close_brace=close_brace,
|
||
|
children=children,
|
||
|
last=root,
|
||
|
)
|
||
|
append = children.append
|
||
|
num_items = len(obj)
|
||
|
last_item_index = num_items - 1
|
||
|
|
||
|
if _safe_isinstance(obj, _MAPPING_CONTAINERS):
|
||
|
iter_items = iter(obj.items())
|
||
|
if max_length is not None:
|
||
|
iter_items = islice(iter_items, max_length)
|
||
|
for index, (key, child) in enumerate(iter_items):
|
||
|
child_node = _traverse(child, depth=depth + 1)
|
||
|
child_node.key_repr = to_repr(key)
|
||
|
child_node.last = index == last_item_index
|
||
|
append(child_node)
|
||
|
else:
|
||
|
iter_values = iter(obj)
|
||
|
if max_length is not None:
|
||
|
iter_values = islice(iter_values, max_length)
|
||
|
for index, child in enumerate(iter_values):
|
||
|
child_node = _traverse(child, depth=depth + 1)
|
||
|
child_node.last = index == last_item_index
|
||
|
append(child_node)
|
||
|
if max_length is not None and num_items > max_length:
|
||
|
append(Node(value_repr=f"... +{num_items - max_length}", last=True))
|
||
|
else:
|
||
|
node = Node(empty=empty, children=[], last=root)
|
||
|
|
||
|
pop_visited(obj_id)
|
||
|
else:
|
||
|
node = Node(value_repr=to_repr(obj), last=root)
|
||
|
node.is_tuple = _safe_isinstance(obj, tuple)
|
||
|
node.is_namedtuple = _is_namedtuple(obj)
|
||
|
return node
|
||
|
|
||
|
node = _traverse(_object, root=True)
|
||
|
return node
|
||
|
|
||
|
|
||
|
def pretty_repr(
|
||
|
_object: Any,
|
||
|
*,
|
||
|
max_width: int = 80,
|
||
|
indent_size: int = 4,
|
||
|
max_length: Optional[int] = None,
|
||
|
max_string: Optional[int] = None,
|
||
|
max_depth: Optional[int] = None,
|
||
|
expand_all: bool = False,
|
||
|
) -> str:
|
||
|
"""Prettify repr string by expanding on to new lines to fit within a given width.
|
||
|
|
||
|
Args:
|
||
|
_object (Any): Object to repr.
|
||
|
max_width (int, optional): Desired maximum width of repr string. Defaults to 80.
|
||
|
indent_size (int, optional): Number of spaces to indent. Defaults to 4.
|
||
|
max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
|
||
|
Defaults to None.
|
||
|
max_string (int, optional): Maximum length of string before truncating, or None to disable truncating.
|
||
|
Defaults to None.
|
||
|
max_depth (int, optional): Maximum depth of nested data structure, or None for no depth.
|
||
|
Defaults to None.
|
||
|
expand_all (bool, optional): Expand all containers regardless of available width. Defaults to False.
|
||
|
|
||
|
Returns:
|
||
|
str: A possibly multi-line representation of the object.
|
||
|
"""
|
||
|
|
||
|
if _safe_isinstance(_object, Node):
|
||
|
node = _object
|
||
|
else:
|
||
|
node = traverse(
|
||
|
_object, max_length=max_length, max_string=max_string, max_depth=max_depth
|
||
|
)
|
||
|
repr_str: str = node.render(
|
||
|
max_width=max_width, indent_size=indent_size, expand_all=expand_all
|
||
|
)
|
||
|
return repr_str
|
||
|
|
||
|
|
||
|
def pprint(
|
||
|
_object: Any,
|
||
|
*,
|
||
|
console: Optional["Console"] = None,
|
||
|
indent_guides: bool = True,
|
||
|
max_length: Optional[int] = None,
|
||
|
max_string: Optional[int] = None,
|
||
|
max_depth: Optional[int] = None,
|
||
|
expand_all: bool = False,
|
||
|
) -> None:
|
||
|
"""A convenience function for pretty printing.
|
||
|
|
||
|
Args:
|
||
|
_object (Any): Object to pretty print.
|
||
|
console (Console, optional): Console instance, or None to use default. Defaults to None.
|
||
|
max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
|
||
|
Defaults to None.
|
||
|
max_string (int, optional): Maximum length of strings before truncating, or None to disable. Defaults to None.
|
||
|
max_depth (int, optional): Maximum depth for nested data structures, or None for unlimited depth. Defaults to None.
|
||
|
indent_guides (bool, optional): Enable indentation guides. Defaults to True.
|
||
|
expand_all (bool, optional): Expand all containers. Defaults to False.
|
||
|
"""
|
||
|
_console = get_console() if console is None else console
|
||
|
_console.print(
|
||
|
Pretty(
|
||
|
_object,
|
||
|
max_length=max_length,
|
||
|
max_string=max_string,
|
||
|
max_depth=max_depth,
|
||
|
indent_guides=indent_guides,
|
||
|
expand_all=expand_all,
|
||
|
overflow="ignore",
|
||
|
),
|
||
|
soft_wrap=True,
|
||
|
)
|
||
|
|
||
|
|
||
|
if __name__ == "__main__": # pragma: no cover
|
||
|
|
||
|
class BrokenRepr:
|
||
|
def __repr__(self) -> str:
|
||
|
1 / 0
|
||
|
return "this will fail"
|
||
|
|
||
|
from typing import NamedTuple
|
||
|
|
||
|
class StockKeepingUnit(NamedTuple):
|
||
|
name: str
|
||
|
description: str
|
||
|
price: float
|
||
|
category: str
|
||
|
reviews: List[str]
|
||
|
|
||
|
d = defaultdict(int)
|
||
|
d["foo"] = 5
|
||
|
data = {
|
||
|
"foo": [
|
||
|
1,
|
||
|
"Hello World!",
|
||
|
100.123,
|
||
|
323.232,
|
||
|
432324.0,
|
||
|
{5, 6, 7, (1, 2, 3, 4), 8},
|
||
|
],
|
||
|
"bar": frozenset({1, 2, 3}),
|
||
|
"defaultdict": defaultdict(
|
||
|
list, {"crumble": ["apple", "rhubarb", "butter", "sugar", "flour"]}
|
||
|
),
|
||
|
"counter": Counter(
|
||
|
[
|
||
|
"apple",
|
||
|
"orange",
|
||
|
"pear",
|
||
|
"kumquat",
|
||
|
"kumquat",
|
||
|
"durian" * 100,
|
||
|
]
|
||
|
),
|
||
|
"atomic": (False, True, None),
|
||
|
"namedtuple": StockKeepingUnit(
|
||
|
"Sparkling British Spring Water",
|
||
|
"Carbonated spring water",
|
||
|
0.9,
|
||
|
"water",
|
||
|
["its amazing!", "its terrible!"],
|
||
|
),
|
||
|
"Broken": BrokenRepr(),
|
||
|
}
|
||
|
data["foo"].append(data) # type: ignore[attr-defined]
|
||
|
|
||
|
from rich import print
|
||
|
|
||
|
print(Pretty(data, indent_guides=True, max_string=20))
|
||
|
|
||
|
class Thing:
|
||
|
def __repr__(self) -> str:
|
||
|
return "Hello\x1b[38;5;239m World!"
|
||
|
|
||
|
print(Pretty(Thing()))
|