Skip to content

Commit af535ce

Browse files
chore(pydantic v1): exclude specific properties when rich printing (#1751)
1 parent 70edb21 commit af535ce

File tree

5 files changed

+33
-9
lines changed

5 files changed

+33
-9
lines changed

src/openai/_models.py

+9-1
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,15 @@
22

33
import os
44
import inspect
5-
from typing import TYPE_CHECKING, Any, Type, Union, Generic, TypeVar, Callable, Optional, cast
5+
from typing import TYPE_CHECKING, Any, Type, Tuple, Union, Generic, TypeVar, Callable, Optional, cast
66
from datetime import date, datetime
77
from typing_extensions import (
88
Unpack,
99
Literal,
1010
ClassVar,
1111
Protocol,
1212
Required,
13+
Sequence,
1314
ParamSpec,
1415
TypedDict,
1516
TypeGuard,
@@ -72,6 +73,8 @@
7273

7374
P = ParamSpec("P")
7475

76+
ReprArgs = Sequence[Tuple[Optional[str], Any]]
77+
7578

7679
@runtime_checkable
7780
class _ConfigProtocol(Protocol):
@@ -94,6 +97,11 @@ def model_fields_set(self) -> set[str]:
9497
class Config(pydantic.BaseConfig): # pyright: ignore[reportDeprecated]
9598
extra: Any = pydantic.Extra.allow # type: ignore
9699

100+
@override
101+
def __repr_args__(self) -> ReprArgs:
102+
# we don't want these attributes to be included when something like `rich.print` is used
103+
return [arg for arg in super().__repr_args__() if arg[0] not in {"_request_id", "__exclude_fields__"}]
104+
97105
if TYPE_CHECKING:
98106
_request_id: Optional[str] = None
99107
"""The ID of the request, returned via the X-Request-ID header. Useful for debugging requests and reporting issues to OpenAI.

tests/lib/chat/_utils.py

+3-8
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,14 @@
11
from __future__ import annotations
22

3-
import io
43
import inspect
54
from typing import Any, Iterable
65
from typing_extensions import TypeAlias
76

8-
import rich
97
import pytest
108
import pydantic
119

10+
from ...utils import rich_print_str
11+
1212
ReprArgs: TypeAlias = "Iterable[tuple[str | None, Any]]"
1313

1414

@@ -26,12 +26,7 @@ def __repr_args__(self: pydantic.BaseModel) -> ReprArgs:
2626
with monkeypatch.context() as m:
2727
m.setattr(pydantic.BaseModel, "__repr_args__", __repr_args__)
2828

29-
buf = io.StringIO()
30-
31-
console = rich.console.Console(file=buf, width=120)
32-
console.print(obj)
33-
34-
string = buf.getvalue()
29+
string = rich_print_str(obj)
3530

3631
# we remove all `fn_name.<locals>.` occurences
3732
# so that we can share the same snapshots between

tests/test_legacy_response.py

+4
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
from openai._base_client import FinalRequestOptions
1212
from openai._legacy_response import LegacyAPIResponse
1313

14+
from .utils import rich_print_str
15+
1416

1517
class PydanticModel(pydantic.BaseModel): ...
1618

@@ -85,6 +87,8 @@ def test_response_basemodel_request_id(client: OpenAI) -> None:
8587
assert obj.foo == "hello!"
8688
assert obj.bar == 2
8789
assert obj.to_dict() == {"foo": "hello!", "bar": 2}
90+
assert "_request_id" not in rich_print_str(obj)
91+
assert "__exclude_fields__" not in rich_print_str(obj)
8892

8993

9094
def test_response_parse_annotated_type(client: OpenAI) -> None:

tests/test_response.py

+4
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
from openai._streaming import Stream
1919
from openai._base_client import FinalRequestOptions
2020

21+
from .utils import rich_print_str
22+
2123

2224
class ConcreteBaseAPIResponse(APIResponse[bytes]): ...
2325

@@ -175,6 +177,8 @@ def test_response_basemodel_request_id(client: OpenAI) -> None:
175177
assert obj.foo == "hello!"
176178
assert obj.bar == 2
177179
assert obj.to_dict() == {"foo": "hello!", "bar": 2}
180+
assert "_request_id" not in rich_print_str(obj)
181+
assert "__exclude_fields__" not in rich_print_str(obj)
178182

179183

180184
@pytest.mark.asyncio

tests/utils.py

+13
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import io
34
import os
45
import inspect
56
import traceback
@@ -8,6 +9,8 @@
89
from datetime import date, datetime
910
from typing_extensions import Literal, get_args, get_origin, assert_type
1011

12+
import rich
13+
1114
from openai._types import Omit, NoneType
1215
from openai._utils import (
1316
is_dict,
@@ -138,6 +141,16 @@ def _assert_list_type(type_: type[object], value: object) -> None:
138141
assert_type(inner_type, entry) # type: ignore
139142

140143

144+
def rich_print_str(obj: object) -> str:
145+
"""Like `rich.print()` but returns the string instead"""
146+
buf = io.StringIO()
147+
148+
console = rich.console.Console(file=buf, width=120)
149+
console.print(obj)
150+
151+
return buf.getvalue()
152+
153+
141154
@contextlib.contextmanager
142155
def update_env(**new_env: str | Omit) -> Iterator[None]:
143156
old = os.environ.copy()

0 commit comments

Comments
 (0)