From 01f95d7e810d0f99cb2406cfdecac841db5ce36b Mon Sep 17 00:00:00 2001 From: Will McGugan Date: Thu, 6 Jan 2022 10:33:18 +0000 Subject: [PATCH 1/5] max depth arg --- rich/pretty.py | 185 +++++++++++++++++++++++++------------------ tests/test_pretty.py | 69 ++++++++++++++++ 2 files changed, 178 insertions(+), 76 deletions(-) diff --git a/rich/pretty.py b/rich/pretty.py index f7059112..80c4df98 100644 --- a/rich/pretty.py +++ b/rich/pretty.py @@ -196,6 +196,7 @@ class Pretty(JupyterMixin): 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. @@ -213,6 +214,7 @@ class Pretty(JupyterMixin): 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, @@ -226,6 +228,7 @@ class Pretty(JupyterMixin): 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 @@ -239,6 +242,7 @@ class Pretty(JupyterMixin): 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( @@ -474,7 +478,10 @@ class _Line: def traverse( - _object: Any, max_length: Optional[int] = None, max_string: Optional[int] = None + _object: Any, + max_length: Optional[int] = None, + max_string: Optional[int] = None, + max_depth: Optional[int] = None, ) -> Node: """Traverse object and generate a tree. @@ -484,6 +491,8 @@ def traverse( 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. @@ -509,11 +518,13 @@ def traverse( push_visited = visited_ids.add pop_visited = visited_ids.remove - def _traverse(obj: Any, root: bool = False) -> Node: + def _traverse(obj: Any, root: bool = False, depth: int = 0) -> Node: """Walk the object depth first.""" + obj_type = type(obj) py_version = (sys.version_info.major, sys.version_info.minor) 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: @@ -554,33 +565,37 @@ def traverse( if args: children = [] append = children.append - if angular: - node = Node( - open_brace=f"<{class_name} ", - close_brace=">", - children=children, - last=root, - separator=" ", - ) + + if reached_max_depth: + node = Node(value_repr=f"...") else: - node = Node( - open_brace=f"{class_name}(", - close_brace=")", - children=children, - last=root, - ) - for last, arg in loop_last(args): - if isinstance(arg, tuple): - key, child = arg - child_node = _traverse(child) - child_node.last = last - child_node.key_repr = key - child_node.key_separator = "=" - append(child_node) + if angular: + node = Node( + open_brace=f"<{class_name} ", + close_brace=">", + children=children, + last=root, + separator=" ", + ) else: - child_node = _traverse(arg) - child_node.last = last - append(child_node) + node = Node( + open_brace=f"{class_name}(", + close_brace=")", + children=children, + last=root, + ) + for last, arg in loop_last(args): + if 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}()", @@ -593,40 +608,43 @@ def traverse( attr_fields = _get_attr_fields(obj) if attr_fields: - node = Node( - open_brace=f"{obj.__class__.__name__}(", - close_brace=")", - children=children, - last=root, - ) + if reached_max_depth: + node = Node(value_repr=f"...") + 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, - ) + 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) - child_node.last = last - child_node.key_repr = name - child_node.key_separator = "=" - append(child_node) + 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 @@ -646,21 +664,26 @@ def traverse( children = [] append = children.append - node = Node( - open_brace=f"{obj.__class__.__name__}(", - close_brace=")", - children=children, - last=root, - ) + if reached_max_depth: + node = Node(value_repr=f"...") + else: + node = Node( + open_brace=f"{obj.__class__.__name__}(", + close_brace=")", + children=children, + last=root, + ) - for last, field in loop_last(field for field in fields(obj) if field.repr): - child_node = _traverse(getattr(obj, field.name)) - child_node.key_repr = field.name - child_node.last = last - child_node.key_separator = "=" - append(child_node) + 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) + pop_visited(obj_id) elif isinstance(obj, _CONTAINERS): for container_type in _CONTAINERS: @@ -676,7 +699,9 @@ def traverse( open_brace, close_brace, empty = _BRACES[obj_type](obj) - if obj_type.__repr__ != type(obj).__repr__: + if reached_max_depth: + node = Node(value_repr=f"...", last=root) + elif obj_type.__repr__ != type(obj).__repr__: node = Node(value_repr=to_repr(obj), last=root) elif obj: children = [] @@ -695,7 +720,7 @@ def traverse( 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) + child_node = _traverse(child, depth=depth + 1) child_node.key_repr = to_repr(key) child_node.last = index == last_item_index append(child_node) @@ -704,7 +729,7 @@ def traverse( if max_length is not None: iter_values = islice(iter_values, max_length) for index, child in enumerate(iter_values): - child_node = _traverse(child) + 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: @@ -729,6 +754,7 @@ def pretty_repr( 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. @@ -741,6 +767,8 @@ def pretty_repr( 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: @@ -750,7 +778,9 @@ def pretty_repr( if isinstance(_object, Node): node = _object else: - node = traverse(_object, max_length=max_length, max_string=max_string) + node = traverse( + _object, max_length=max_length, max_string=max_string, max_depth=max_depth + ) repr_str = node.render( max_width=max_width, indent_size=indent_size, expand_all=expand_all ) @@ -764,6 +794,7 @@ def pprint( 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. @@ -774,6 +805,7 @@ def pprint( 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 no maximum. Defaults to None. indent_guides (bool, optional): Enable indentation guides. Defaults to True. expand_all (bool, optional): Expand all containers. Defaults to False. """ @@ -783,6 +815,7 @@ def pprint( _object, max_length=max_length, max_string=max_string, + max_depth=max_depth, indent_guides=indent_guides, expand_all=expand_all, overflow="ignore", diff --git a/tests/test_pretty.py b/tests/test_pretty.py index ab83d930..1752d1e6 100644 --- a/tests/test_pretty.py +++ b/tests/test_pretty.py @@ -131,6 +131,75 @@ def test_recursive(): assert result == expected +def test_max_depth(): + d = d = {} + d["foo"] = {"fob": {"a": [1, 2, 3], "b": {"z": "x", "y": ["a", "b", "c"]}}} + + assert pretty_repr(d, max_depth=0) == "..." + assert pretty_repr(d, max_depth=1) == "{'foo': ...}" + assert pretty_repr(d, max_depth=2) == "{'foo': {'fob': ...}}" + assert pretty_repr(d, max_depth=3) == "{'foo': {'fob': {'a': ..., 'b': ...}}}" + assert ( + pretty_repr(d, max_width=100, max_depth=4) + == "{'foo': {'fob': {'a': [1, 2, 3], 'b': {'z': 'x', 'y': ...}}}}" + ) + assert ( + pretty_repr(d, max_width=100, max_depth=5) + == "{'foo': {'fob': {'a': [1, 2, 3], 'b': {'z': 'x', 'y': ['a', 'b', 'c']}}}}" + ) + + +def test_max_depth_rich_repr(): + class Foo: + def __init__(self, foo): + self.foo = foo + + def __rich_repr__(self): + yield "foo", self.foo + + class Bar: + def __init__(self, bar): + self.bar = bar + + def __rich_repr__(self): + yield "bar", self.bar + + assert ( + pretty_repr(Foo(foo=Bar(bar=Foo(foo=[]))), max_depth=2) + == "Foo(foo=Bar(bar=...))" + ) + + +def test_max_depth_attrs(): + @attr.define + class Foo: + foo = attr.field() + + @attr.define + class Bar: + bar = attr.field() + + assert ( + pretty_repr(Foo(foo=Bar(bar=Foo(foo=[]))), max_depth=2) + == "Foo(foo=Bar(bar=...))" + ) + + +def test_max_depth_dataclass(): + @dataclass + class Foo: + foo: object + + @dataclass + class Bar: + bar: object + + assert ( + pretty_repr(Foo(foo=Bar(bar=Foo(foo=[]))), max_depth=2) + == "Foo(foo=Bar(bar=...))" + ) + + def test_defaultdict(): test_dict = defaultdict(int, {"foo": 2}) result = pretty_repr(test_dict) From 3999c8a8e696eb853ba00b1b49c4fa049170ad33 Mon Sep 17 00:00:00 2001 From: Will McGugan Date: Thu, 6 Jan 2022 10:45:33 +0000 Subject: [PATCH 2/5] changelog --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index db2fb521..a5f50c90 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [11.0.0] - Unreleased + +### Added + +- Added max_depth arg to pretty printing + ## [10.16.2] - 2021-01-02 ### Fixed From 76b5dd233a866f8975e15fa92849e1ff743b6067 Mon Sep 17 00:00:00 2001 From: Will McGugan Date: Thu, 6 Jan 2022 15:04:47 +0000 Subject: [PATCH 3/5] Update tests/test_pretty.py Co-authored-by: Darren Burns --- tests/test_pretty.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_pretty.py b/tests/test_pretty.py index 1752d1e6..40bb485f 100644 --- a/tests/test_pretty.py +++ b/tests/test_pretty.py @@ -132,7 +132,7 @@ def test_recursive(): def test_max_depth(): - d = d = {} + d = {} d["foo"] = {"fob": {"a": [1, 2, 3], "b": {"z": "x", "y": ["a", "b", "c"]}}} assert pretty_repr(d, max_depth=0) == "..." From 3974a9ef828e724ba57f34ca977962325a3bedae Mon Sep 17 00:00:00 2001 From: Will McGugan Date: Thu, 6 Jan 2022 15:05:11 +0000 Subject: [PATCH 4/5] Update rich/pretty.py Co-authored-by: Darren Burns --- rich/pretty.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rich/pretty.py b/rich/pretty.py index 80c4df98..10cc2925 100644 --- a/rich/pretty.py +++ b/rich/pretty.py @@ -805,7 +805,7 @@ def pprint( 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 no maximum. 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. """ From bff7981788dd9a29bff771f0bc2267d1e3936f31 Mon Sep 17 00:00:00 2001 From: Will McGugan Date: Thu, 6 Jan 2022 15:07:05 +0000 Subject: [PATCH 5/5] test max depth of None --- tests/test_pretty.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_pretty.py b/tests/test_pretty.py index 1752d1e6..e955bd61 100644 --- a/tests/test_pretty.py +++ b/tests/test_pretty.py @@ -147,6 +147,10 @@ def test_max_depth(): pretty_repr(d, max_width=100, max_depth=5) == "{'foo': {'fob': {'a': [1, 2, 3], 'b': {'z': 'x', 'y': ['a', 'b', 'c']}}}}" ) + assert ( + pretty_repr(d, max_width=100, max_depth=None) + == "{'foo': {'fob': {'a': [1, 2, 3], 'b': {'z': 'x', 'y': ['a', 'b', 'c']}}}}" + ) def test_max_depth_rich_repr():