diff --git a/devtools/prettier.py b/devtools/prettier.py index b603bef..6f47ecc 100644 --- a/devtools/prettier.py +++ b/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 diff --git a/tests/test_prettier.py b/tests/test_prettier.py index cf069a9..b55e8fb 100644 --- a/tests/test_prettier.py +++ b/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)