Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore(internal): support parsing Annotated types #1222

Merged
merged 1 commit into from
Mar 8, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
11 changes: 10 additions & 1 deletion src/openai/_legacy_response.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
import pydantic

from ._types import NoneType
from ._utils import is_given
from ._utils import is_given, extract_type_arg, is_annotated_type
from ._models import BaseModel, is_basemodel
from ._constants import RAW_RESPONSE_HEADER
from ._streaming import Stream, AsyncStream, is_stream_class_type, extract_stream_chunk_type
Expand Down Expand Up @@ -174,6 +174,10 @@ def elapsed(self) -> datetime.timedelta:
return self.http_response.elapsed

def _parse(self, *, to: type[_T] | None = None) -> R | _T:
# unwrap `Annotated[T, ...]` -> `T`
if to and is_annotated_type(to):
to = extract_type_arg(to, 0)

if self._stream:
if to:
if not is_stream_class_type(to):
Expand Down Expand Up @@ -215,6 +219,11 @@ def _parse(self, *, to: type[_T] | None = None) -> R | _T:
)

cast_to = to if to is not None else self._cast_to

# unwrap `Annotated[T, ...]` -> `T`
if is_annotated_type(cast_to):
cast_to = extract_type_arg(cast_to, 0)

if cast_to is NoneType:
return cast(R, None)

Expand Down
14 changes: 13 additions & 1 deletion src/openai/_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,16 @@
AnyMapping,
HttpxRequestFiles,
)
from ._utils import is_list, is_given, is_mapping, parse_date, parse_datetime, strip_not_given
from ._utils import (
is_list,
is_given,
is_mapping,
parse_date,
parse_datetime,
strip_not_given,
extract_type_arg,
is_annotated_type,
)
from ._compat import (
PYDANTIC_V2,
ConfigDict,
Expand Down Expand Up @@ -275,6 +284,9 @@ def construct_type(*, value: object, type_: type) -> object:

If the given value does not match the expected type then it is returned as-is.
"""
# unwrap `Annotated[T, ...]` -> `T`
if is_annotated_type(type_):
type_ = extract_type_arg(type_, 0)

# we need to use the origin class for any types that are subscripted generics
# e.g. Dict[str, object]
Expand Down
11 changes: 10 additions & 1 deletion src/openai/_response.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
import pydantic

from ._types import NoneType
from ._utils import is_given, extract_type_var_from_base
from ._utils import is_given, extract_type_arg, is_annotated_type, extract_type_var_from_base
from ._models import BaseModel, is_basemodel
from ._constants import RAW_RESPONSE_HEADER, OVERRIDE_CAST_TO_HEADER
from ._streaming import Stream, AsyncStream, is_stream_class_type, extract_stream_chunk_type
Expand Down Expand Up @@ -121,6 +121,10 @@ def __repr__(self) -> str:
)

def _parse(self, *, to: type[_T] | None = None) -> R | _T:
# unwrap `Annotated[T, ...]` -> `T`
if to and is_annotated_type(to):
to = extract_type_arg(to, 0)

if self._is_sse_stream:
if to:
if not is_stream_class_type(to):
Expand Down Expand Up @@ -162,6 +166,11 @@ def _parse(self, *, to: type[_T] | None = None) -> R | _T:
)

cast_to = to if to is not None else self._cast_to

# unwrap `Annotated[T, ...]` -> `T`
if is_annotated_type(cast_to):
cast_to = extract_type_arg(cast_to, 0)

if cast_to is NoneType:
return cast(R, None)

Expand Down
19 changes: 19 additions & 0 deletions tests/test_legacy_response.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import json
from typing import cast
from typing_extensions import Annotated

import httpx
import pytest
Expand Down Expand Up @@ -63,3 +65,20 @@ def test_response_parse_custom_model(client: OpenAI) -> None:
obj = response.parse(to=CustomModel)
assert obj.foo == "hello!"
assert obj.bar == 2


def test_response_parse_annotated_type(client: OpenAI) -> None:
response = LegacyAPIResponse(
raw=httpx.Response(200, content=json.dumps({"foo": "hello!", "bar": 2})),
client=client,
stream=False,
stream_cls=None,
cast_to=str,
options=FinalRequestOptions.construct(method="get", url="/foo"),
)

obj = response.parse(
to=cast("type[CustomModel]", Annotated[CustomModel, "random metadata"]),
)
assert obj.foo == "hello!"
assert obj.bar == 2
16 changes: 14 additions & 2 deletions tests/test_models.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,14 @@
import json
from typing import Any, Dict, List, Union, Optional, cast
from datetime import datetime, timezone
from typing_extensions import Literal
from typing_extensions import Literal, Annotated

import pytest
import pydantic
from pydantic import Field

from openai._compat import PYDANTIC_V2, parse_obj, model_dump, model_json
from openai._models import BaseModel
from openai._models import BaseModel, construct_type


class BasicModel(BaseModel):
Expand Down Expand Up @@ -571,3 +571,15 @@ class OurModel(BaseModel):
foo: Optional[str] = None

takes_pydantic(OurModel())


def test_annotated_types() -> None:
class Model(BaseModel):
value: str

m = construct_type(
value={"value": "foo"},
type_=cast(Any, Annotated[Model, "random metadata"]),
)
assert isinstance(m, Model)
assert m.value == "foo"
37 changes: 36 additions & 1 deletion tests/test_response.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import json
from typing import List
from typing import List, cast
from typing_extensions import Annotated

import httpx
import pytest
Expand Down Expand Up @@ -157,3 +158,37 @@ async def test_async_response_parse_custom_model(async_client: AsyncOpenAI) -> N
obj = await response.parse(to=CustomModel)
assert obj.foo == "hello!"
assert obj.bar == 2


def test_response_parse_annotated_type(client: OpenAI) -> None:
response = APIResponse(
raw=httpx.Response(200, content=json.dumps({"foo": "hello!", "bar": 2})),
client=client,
stream=False,
stream_cls=None,
cast_to=str,
options=FinalRequestOptions.construct(method="get", url="/foo"),
)

obj = response.parse(
to=cast("type[CustomModel]", Annotated[CustomModel, "random metadata"]),
)
assert obj.foo == "hello!"
assert obj.bar == 2


async def test_async_response_parse_annotated_type(async_client: AsyncOpenAI) -> None:
response = AsyncAPIResponse(
raw=httpx.Response(200, content=json.dumps({"foo": "hello!", "bar": 2})),
client=async_client,
stream=False,
stream_cls=None,
cast_to=str,
options=FinalRequestOptions.construct(method="get", url="/foo"),
)

obj = await response.parse(
to=cast("type[CustomModel]", Annotated[CustomModel, "random metadata"]),
)
assert obj.foo == "hello!"
assert obj.bar == 2
6 changes: 6 additions & 0 deletions tests/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
is_list,
is_list_type,
is_union_type,
extract_type_arg,
is_annotated_type,
)
from openai._compat import PYDANTIC_V2, field_outer_type, get_model_fields
from openai._models import BaseModel
Expand Down Expand Up @@ -49,6 +51,10 @@ def assert_matches_type(
path: list[str],
allow_none: bool = False,
) -> None:
# unwrap `Annotated[T, ...]` -> `T`
if is_annotated_type(type_):
type_ = extract_type_arg(type_, 0)

if allow_none and value is None:
return

Expand Down