Repository: Textualize/rich
Pull request #2524 (2022-09-14 20:52:03)

Fix `pretty` cyclic reference handling
## Type of changes

- [x] Bug fix
- [ ] New feature
- [ ] Documentation / docstrings
- [ ] Tests
- [ ] Other

## Checklist

- [x] I've run the latest [black](https://github.com/psf/black) with default args on new code.
- [x] I've updated CHANGELOG.md and CONTRIBUTORS.md where appropriate.
- [x] I've added tests for new code.
- [x] I accept that @willmcgugan may be pedantic in the code review.

## Description

Previously cyclic references were only handled for container and dataclass types, but not namedtuple, attrs, or custom types. This fixes that, and expands the tests to cover these cases.



Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 409c0280..d678c005 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,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
 - Fix missing `mode` property on file wrapper breaking uploads via `requests` https://github.com/Textualize/rich/pull/2495
 
 ## [12.5.2] - 2022-07-18
diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
index 533eb2bf..1972229a 100644
--- a/CONTRIBUTORS.md
+++ b/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)
diff --git a/rich/pretty.py b/rich/pretty.py
index 12c2454d..3e83b200 100644
--- a/rich/pretty.py
+++ b/rich/pretty.py
@@ -636,6 +636,11 @@ def traverse(
     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 traverse(
                 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 traverse(
                     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 traverse(
                 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 traverse(
                     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 traverse(
                     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)


Test patch:
diff --git a/tests/test_pretty.py b/tests/test_pretty.py
index 163a9d92..2b24d907 100644
--- a/tests/test_pretty.py
+++ b/tests/test_pretty.py
@@ -4,7 +4,7 @@ import sys
 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 test_broken_getattr():
     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():