openai/openai-python
Publicmirrored from https://github.com/openai/openai-pythonAvailable
tests/lib/chat/_utils.py
59lines · modeblame
bf1ca86cRobert Craigie1 years ago | 1 | from __future__ import annotations |
| 2 | | |
| 3 | import io | |
| 4 | import inspect | |
| 5 | from typing import Any, Iterable | |
| 6 | from typing_extensions import TypeAlias | |
| 7 | | |
| 8 | import rich | |
| 9 | import pytest | |
| 10 | import pydantic | |
| 11 | | |
| 12 | ReprArgs: TypeAlias = "Iterable[tuple[str | None, Any]]" | |
| 13 | | |
| 14 | | |
| 15 | def print_obj(obj: object, monkeypatch: pytest.MonkeyPatch) -> str: | |
| 16 | """Pretty print an object to a string""" | |
| 17 | | |
| 18 | # monkeypatch pydantic model printing so that model fields | |
| 19 | # are always printed in the same order so we can reliably | |
| 20 | # use this for snapshot tests | |
| 21 | original_repr = pydantic.BaseModel.__repr_args__ | |
| 22 | | |
| 23 | def __repr_args__(self: pydantic.BaseModel) -> ReprArgs: | |
| 24 | return sorted(original_repr(self), key=lambda arg: arg[0] or arg) | |
| 25 | | |
| 26 | with monkeypatch.context() as m: | |
| 27 | m.setattr(pydantic.BaseModel, "__repr_args__", __repr_args__) | |
| 28 | | |
| 29 | buf = io.StringIO() | |
| 30 | | |
| 31 | console = rich.console.Console(file=buf, width=120) | |
| 32 | console.print(obj) | |
| 33 | | |
| 34 | string = buf.getvalue() | |
| 35 | | |
| 36 | # we remove all `fn_name.<locals>.` occurences | |
| 37 | # so that we can share the same snapshots between | |
| 38 | # pydantic v1 and pydantic v2 as their output for | |
| 39 | # generic models differs, e.g. | |
| 40 | # | |
| 41 | # v2: `ParsedChatCompletion[test_parse_pydantic_model.<locals>.Location]` | |
| 42 | # v1: `ParsedChatCompletion[Location]` | |
| 43 | return clear_locals(string, stacklevel=2) | |
| 44 | | |
| 45 | | |
| 46 | def get_caller_name(*, stacklevel: int = 1) -> str: | |
| 47 | frame = inspect.currentframe() | |
| 48 | assert frame is not None | |
| 49 | | |
| 50 | for i in range(stacklevel): | |
| 51 | frame = frame.f_back | |
| 52 | assert frame is not None, f"no {i}th frame" | |
| 53 | | |
| 54 | return frame.f_code.co_name | |
| 55 | | |
| 56 | | |
| 57 | def clear_locals(string: str, *, stacklevel: int) -> str: | |
| 58 | caller = get_caller_name(stacklevel=stacklevel + 1) | |
| 59 | return string.replace(f"{caller}.<locals>.", "") |