Restore PrettyFormat streams after reentrant calls (#169)

Co-authored-by: Samuel Colvin <s@muelcolvin.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Author: 3 people

devtools/prettier.py

@@ -66,6 +66,7 @@ def __init__(
         width: int = 120,
         yield_from_generators: bool = True,
     ):
+        self._stream = io.StringIO()
         self._indent_step = indent_step
         self._c = indent_char
         self._repr_strings = repr_strings
@@ -87,9 +88,15 @@ def __init__(
         ]
 
     def __call__(self, value: Any, *, indent: int = 0, indent_first: bool = False, highlight: bool = False) -> str:
+        # save and restore the stream so `pformat` can be called from inside a `__repr__` without clobbering
+        # the output of the outer call, see #162
+        previous_stream = self._stream
         self._stream = io.StringIO()
-        self._format(value, indent_current=indent, indent_first=indent_first)
-        s = self._stream.getvalue()
+        try:
+            self._format(value, indent_current=indent, indent_first=indent_first)
+            s = self._stream.getvalue()
+        finally:
+            self._stream = previous_stream
         pyg = get_pygments()
         if highlight and pyg:
             pygments, pyg_lexer, pyg_formatter = pyg

tests/test_prettier.py

@@ -33,6 +33,53 @@ def test_print(capsys: pytest.CaptureFixture[str]):
     assert stderr == ''
 
 
+@pytest.mark.parametrize('depth', [1, 3])
+@pytest.mark.parametrize('container', [False, True])
+def test_reentrant_repr(depth: int, container: bool):
+    formatter = PrettyFormat(simple_cutoff=0)
+
+    class Reentrant:
+        def __init__(self, depth: int) -> None:
+            self.depth = depth
+
+        def __repr__(self) -> str:
+            value = Reentrant(self.depth - 1) if self.depth > 1 else 1
+            return formatter(value)
+
+    value = [Reentrant(depth), 2] if container else Reentrant(depth)
+    expected = '[\n    1,\n    2,\n]' if container else '1'
+    assert formatter(value) == expected
+    assert formatter(3) == '3'
+
+
+def test_reentrant_pprint(capsys: pytest.CaptureFixture[str]):
+    class Reentrant:
+        def __repr__(self) -> str:
+            return pformat(1)
+
+    pprint(Reentrant())
+    assert strip_ansi(capsys.readouterr().out) == '1\n'
+
+
+def test_reentrant_exception():
+    formatter = PrettyFormat(simple_cutoff=0)
+
+    class Broken:
+        def __repr__(self) -> str:
+            raise ValueError('broken repr')
+
+    class Reentrant:
+        def __repr__(self) -> str:
+            with pytest.raises(ValueError, match='broken repr'):
+                formatter(Broken())
+            return '1'
+
+    assert formatter([Reentrant(), 2]) == '[\n    1,\n    2,\n]'
+    with pytest.raises(ValueError, match='broken repr'):
+        formatter(Broken())
+    assert formatter([1, 2]) == '[\n    1,\n    2,\n]'
+
+
 def test_colours():
     v = pformat({1: 2, 3: 4}, highlight=True)
     assert v.startswith('\x1b'), repr(v)