Fix pretty cyclic reference handling

CHANGELOG.md

@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Handle stdout/stderr being null https://github.com/Textualize/rich/pull/2513
 - Fix NO_COLOR support on legacy Windows https://github.com/Textualize/rich/pull/2458
+- Fix pretty printer handling of cyclic references https://github.com/Textualize/rich/pull/2524
 
 ## [12.5.2] - 2022-07-18
 

CONTRIBUTORS.md

@@ -8,6 +8,7 @@ The following people have contributed to the development of Rich:
 - [Gregory Beauregard](https://github.com/GBeauregard/pyffstream)
 - [Dennis Brakhane](https://github.com/brakhane)
 - [Darren Burns](https://github.com/darrenburns)
+- [Jim Crist-Harif](https://github.com/jcrist)
 - [Ed Davis](https://github.com/davised)
 - [Pete Davison](https://github.com/pd93)
 - [James Estevez](https://github.com/jstvz)

rich/pretty.py

@@ -636,6 +636,11 @@ def to_repr(obj: Any) -> str:
     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)
         py_version = (sys.version_info.major, sys.version_info.minor)
         children: List[Node]
@@ -673,6 +678,7 @@ def iter_rich_args(rich_args: Any) -> Iterable[Union[Any, Tuple[str, Any]]]:
                 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__
@@ -720,7 +726,9 @@ def iter_rich_args(rich_args: Any) -> Iterable[Union[Any, Tuple[str, Any]]]:
                     children=[],
                     last=root,
                 )
+            pop_visited(obj_id)
         elif _is_attr_object(obj) and not fake_attributes:
+            push_visited(obj_id)
             children = []
             append = children.append
 
@@ -767,19 +775,14 @@ def iter_attrs() -> Iterable[
                 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) or py_version == (3, 6))
         ):
-            obj_id = id(obj)
-            if obj_id in visited_ids:
-                # Recursion detected
-                return Node(value_repr="...")
             push_visited(obj_id)
-
             children = []
             append = children.append
             if reached_max_depth:
@@ -801,8 +804,9 @@ def iter_attrs() -> Iterable[
                     child_node.key_separator = "="
                     append(child_node)
 
-                pop_visited(obj_id)
+            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
@@ -824,16 +828,13 @@ def iter_attrs() -> Iterable[
                     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
 
-            obj_id = id(obj)
-            if obj_id in visited_ids:
-                # Recursion detected
-                return Node(value_repr="...")
             push_visited(obj_id)
 
             open_brace, close_brace, empty = _BRACES[obj_type](obj)

tests/test_pretty.py

@@ -4,7 +4,7 @@
 from array import array
 from collections import UserDict, defaultdict
 from dataclasses import dataclass, field
-from typing import List, NamedTuple
+from typing import List, NamedTuple, Any
 from unittest.mock import patch
 
 import attr
@@ -315,12 +315,112 @@ def __repr__(self):
     assert result == "BrokenAttr()"
 
 
-def test_recursive():
+def test_reference_cycle_container():
     test = []
     test.append(test)
-    result = pretty_repr(test)
-    expected = "[...]"
-    assert result == expected
+    res = pretty_repr(test)
+    assert res == "[...]"
+
+    test = [1, []]
+    test[1].append(test)
+    res = pretty_repr(test)
+    assert res == "[1, [...]]"
+
+    # Not a cyclic reference, just a repeated reference
+    a = [2]
+    test = [1, [a, a]]
+    res = pretty_repr(test)
+    assert res == "[1, [[2], [2]]]"
+
+
+def test_reference_cycle_namedtuple():
+    class Example(NamedTuple):
+        x: int
+        y: Any
+
+    test = Example(1, [Example(2, [])])
+    test.y[0].y.append(test)
+    res = pretty_repr(test)
+    assert res == "Example(x=1, y=[Example(x=2, y=[...])])"
+
+    # Not a cyclic reference, just a repeated reference
+    a = Example(2, None)
+    test = Example(1, [a, a])
+    res = pretty_repr(test)
+    assert res == "Example(x=1, y=[Example(x=2, y=None), Example(x=2, y=None)])"
+
+
+def test_reference_cycle_dataclass():
+    @dataclass
+    class Example:
+        x: int
+        y: Any
+
+    test = Example(1, None)
+    test.y = test
+    res = pretty_repr(test)
+    assert res == "Example(x=1, y=...)"
+
+    test = Example(1, Example(2, None))
+    test.y.y = test
+    res = pretty_repr(test)
+    assert res == "Example(x=1, y=Example(x=2, y=...))"
+
+    # Not a cyclic reference, just a repeated reference
+    a = Example(2, None)
+    test = Example(1, [a, a])
+    res = pretty_repr(test)
+    assert res == "Example(x=1, y=[Example(x=2, y=None), Example(x=2, y=None)])"
+
+
+def test_reference_cycle_attrs():
+    @attr.define
+    class Example:
+        x: int
+        y: Any
+
+    test = Example(1, None)
+    test.y = test
+    res = pretty_repr(test)
+    assert res == "Example(x=1, y=...)"
+
+    test = Example(1, Example(2, None))
+    test.y.y = test
+    res = pretty_repr(test)
+    assert res == "Example(x=1, y=Example(x=2, y=...))"
+
+    # Not a cyclic reference, just a repeated reference
+    a = Example(2, None)
+    test = Example(1, [a, a])
+    res = pretty_repr(test)
+    assert res == "Example(x=1, y=[Example(x=2, y=None), Example(x=2, y=None)])"
+
+
+def test_reference_cycle_custom_repr():
+    class Example:
+        def __init__(self, x, y):
+            self.x = x
+            self.y = y
+
+        def __rich_repr__(self):
+            yield ("x", self.x)
+            yield ("y", self.y)
+
+    test = Example(1, None)
+    test.y = test
+    res = pretty_repr(test)
+    assert res == "Example(x=1, y=...)"
+
+    test = Example(1, Example(2, None))
+    test.y.y = test
+    res = pretty_repr(test)
+    assert res == "Example(x=1, y=Example(x=2, y=...))"
+
+    # Not a cyclic reference, just a repeated reference
+    a = Example(2, None)
+    test = Example(1, [a, a])
+    res = pretty_repr(test)
+    assert res == "Example(x=1, y=[Example(x=2, y=None), Example(x=2, y=None)])"
 
 
 def test_max_depth():