openai/openai-python
Publicmirrored from https://github.com/openai/openai-pythonAvailable
tests/lib/chat/_utils.py
54lines · modeblame
bf1ca86cRobert Craigie1 years ago | 1 | from __future__ import annotations |
| 2 | | |
| 3 | import inspect | |
| 4 | from typing import Any, Iterable | |
| 5 | from typing_extensions import TypeAlias | |
| 6 | | |
| 7 | import pytest | |
| 8 | import pydantic | |
| 9 | | |
e1aeeb0eRobert Craigie1 years ago | 10 | from ...utils import rich_print_str |
| 11 | | |
bf1ca86cRobert Craigie1 years ago | 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 | | |
e1aeeb0eRobert Craigie1 years ago | 29 | string = rich_print_str(obj) |
bf1ca86cRobert Craigie1 years ago | 30 | |
| 31 | # we remove all `fn_name.<locals>.` occurences | |
| 32 | # so that we can share the same snapshots between | |
| 33 | # pydantic v1 and pydantic v2 as their output for | |
| 34 | # generic models differs, e.g. | |
| 35 | # | |
| 36 | # v2: `ParsedChatCompletion[test_parse_pydantic_model.<locals>.Location]` | |
| 37 | # v1: `ParsedChatCompletion[Location]` | |
| 38 | return clear_locals(string, stacklevel=2) | |
| 39 | | |
| 40 | | |
| 41 | def get_caller_name(*, stacklevel: int = 1) -> str: | |
| 42 | frame = inspect.currentframe() | |
| 43 | assert frame is not None | |
| 44 | | |
| 45 | for i in range(stacklevel): | |
| 46 | frame = frame.f_back | |
| 47 | assert frame is not None, f"no {i}th frame" | |
| 48 | | |
| 49 | return frame.f_code.co_name | |
| 50 | | |
| 51 | | |
| 52 | def clear_locals(string: str, *, stacklevel: int) -> str: | |
| 53 | caller = get_caller_name(stacklevel=stacklevel + 1) | |
| 54 | return string.replace(f"{caller}.<locals>.", "") |