Isolate concurrent pretty-format streams

devtools/prettier.py

@@ -3,6 +3,7 @@
 import os
 from collections import OrderedDict
 from collections.abc import Callable, Generator, Iterable, Mapping
+from contextvars import ContextVar
 from functools import cache
 from typing import TYPE_CHECKING, Any, TextIO, cast
 
@@ -27,6 +28,7 @@
 DEFAULT_WIDTH = int(os.getenv('PY_DEVTOOLS_WIDTH', '120'))
 MISSING = object()
 PRETTY_KEY = '__prettier_formatted_value__'
+_current_stream: ContextVar[io.StringIO] = ContextVar('devtools_pretty_stream')
 
 
 def fmt(v: Any) -> Any:
@@ -87,16 +89,23 @@ def __init__(
             (SQLAlchemyClassType, self._format_sqlalchemy_class),
         ]
 
+    @property
+    def _stream(self) -> io.StringIO:
+        return _current_stream.get(self.__dict__['_stream'])
+
+    @_stream.setter
+    def _stream(self, stream: io.StringIO) -> None:
+        # Keep the idle stream in instance state for copying and pickling.
+        self.__dict__['_stream'] = stream
+
     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()
+        # Each call owns its stream, including reentrant and concurrent calls.
+        stream_token = _current_stream.set(io.StringIO())
         try:
             self._format(value, indent_current=indent, indent_first=indent_first)
             s = self._stream.getvalue()
         finally:
-            self._stream = previous_stream
+            _current_stream.reset(stream_token)
         pyg = get_pygments()
         if highlight and pyg:
             pygments, pyg_lexer, pyg_formatter = pyg

tests/test_prettier.py

@@ -5,6 +5,7 @@
 from collections.abc import Iterator
 from dataclasses import dataclass
 from importlib.util import find_spec
+from threading import Event, Thread
 from typing import Any, ClassVar, NamedTuple, cast
 from unittest.mock import MagicMock
 
@@ -61,6 +62,86 @@ def __repr__(self) -> str:
     assert strip_ansi(capsys.readouterr().out) == '1\n'
 
 
+@pytest.mark.parametrize('shared_global', [False, True])
+def test_concurrent_formatter_streams(shared_global: bool):
+    formatter = pformat if shared_global else PrettyFormat()
+    first_entered = Event()
+    second_entered = Event()
+    first_finished = Event()
+    results: dict[str, str] = {}
+
+    class First:
+        def __pretty__(self, fmt: Any, skip_exc: Any) -> Iterator[str]:
+            yield 'First['
+            first_entered.set()
+            assert second_entered.wait(5)
+            yield '1]'
+
+    class Second:
+        def __pretty__(self, fmt: Any, skip_exc: Any) -> Iterator[str]:
+            yield 'Second['
+            second_entered.set()
+            assert first_finished.wait(5)
+            yield '2]'
+
+    def format_first() -> None:
+        results['first'] = formatter(First())
+        first_finished.set()
+
+    def format_second() -> None:
+        assert first_entered.wait(5)
+        results['second'] = formatter(Second())
+
+    threads = [Thread(target=format_first), Thread(target=format_second)]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join(timeout=10)
+        assert not thread.is_alive()
+
+    assert results == {'first': 'First[1]', 'second': 'Second[2]'}
+    assert formatter(3) == '3'
+
+
+def test_different_formatter_inside_pretty():
+    class Nested:
+        def __pretty__(self, fmt: Any, skip_exc: Any) -> Iterator[str]:
+            yield 'Nested('
+            yield PrettyFormat(indent_step=2)([1, 2, 3, 4])
+            yield ')'
+
+    assert pformat(Nested()) == 'Nested([\n  1,\n  2,\n  3,\n  4,\n])'
+
+
+def test_recursive_global_formatter_inside_pretty():
+    class Nested:
+        def __pretty__(self, fmt: Any, skip_exc: Any) -> Iterator[str]:
+            yield 'Nested('
+            yield pformat([1, 2, 3, 4])
+            yield ')'
+
+    assert pformat(Nested()) == 'Nested([\n    1,\n    2,\n    3,\n    4,\n])'
+
+
+def test_different_formatter_exception_inside_pretty():
+    class Broken:
+        def __repr__(self) -> str:
+            raise ValueError('broken repr')
+
+    class Nested:
+        def __pretty__(self, fmt: Any, skip_exc: Any) -> Iterator[str]:
+            yield 'before:'
+            with pytest.raises(ValueError, match='broken repr'):
+                PrettyFormat()(Broken())
+            yield pformat(1)
+            yield ':after'
+
+    assert pformat(Nested()) == 'before:1:after'
+    with pytest.raises(ValueError, match='broken repr'):
+        pformat(Broken())
+    assert pformat(2) == '2'
+
+
 def test_reentrant_exception():
     formatter = PrettyFormat(simple_cutoff=0)