openai/openai-python
Publicmirrored fromhttps://github.com/openai/openai-pythonAvailable
tests/lib/chat/_utils.py
59lines · modecode
| 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>.", "") |
| 60 | |