avoid misparsing references as other types using Pydantic discriminated unions

openapi_python_client/schema/openapi_schema_pydantic/components.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional
 
 from pydantic import BaseModel, ConfigDict
 
@@ -7,7 +7,7 @@
 from .header import Header
 from .link import Link
 from .parameter import Parameter
-from .reference import Reference
+from .reference import ReferenceOr
 from .request_body import RequestBody
 from .response import Response
 from .schema import Schema
@@ -25,15 +25,15 @@ class Components(BaseModel):
         - https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.3.md#componentsObject
     """
 
-    schemas: Optional[dict[str, Union[Schema, Reference]]] = None
-    responses: Optional[dict[str, Union[Response, Reference]]] = None
-    parameters: Optional[dict[str, Union[Parameter, Reference]]] = None
-    examples: Optional[dict[str, Union[Example, Reference]]] = None
-    requestBodies: Optional[dict[str, Union[RequestBody, Reference]]] = None
-    headers: Optional[dict[str, Union[Header, Reference]]] = None
-    securitySchemes: Optional[dict[str, Union[SecurityScheme, Reference]]] = None
-    links: Optional[dict[str, Union[Link, Reference]]] = None
-    callbacks: Optional[dict[str, Union[Callback, Reference]]] = None
+    schemas: Optional[dict[str, ReferenceOr[Schema]]] = None
+    responses: Optional[dict[str, ReferenceOr[Response]]] = None
+    parameters: Optional[dict[str, ReferenceOr[Parameter]]] = None
+    examples: Optional[dict[str, ReferenceOr[Example]]] = None
+    requestBodies: Optional[dict[str, ReferenceOr[RequestBody]]] = None
+    headers: Optional[dict[str, ReferenceOr[Header]]] = None
+    securitySchemes: Optional[dict[str, ReferenceOr[SecurityScheme]]] = None
+    links: Optional[dict[str, ReferenceOr[Link]]] = None
+    callbacks: Optional[dict[str, ReferenceOr[Callback]]] = None
     model_config = ConfigDict(
         # `Callback` contains an unresolvable forward reference, will rebuild in `__init__.py`:
         defer_build=True,

openapi_python_client/schema/openapi_schema_pydantic/encoding.py

@@ -1,8 +1,8 @@
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Optional
 
 from pydantic import BaseModel, ConfigDict
 
-from .reference import Reference
+from .reference import ReferenceOr
 
 if TYPE_CHECKING:  # pragma: no cover
     from .header import Header
@@ -17,7 +17,7 @@ class Encoding(BaseModel):
     """
 
     contentType: Optional[str] = None
-    headers: Optional[dict[str, Union["Header", Reference]]] = None
+    headers: Optional[dict[str, ReferenceOr["Header"]]] = None
     style: Optional[str] = None
     explode: bool = False
     allowReserved: bool = False

openapi_python_client/schema/openapi_schema_pydantic/media_type.py

@@ -1,10 +1,10 @@
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 from pydantic import BaseModel, ConfigDict, Field
 
 from .encoding import Encoding
 from .example import Example
-from .reference import Reference
+from .reference import ReferenceOr
 from .schema import Schema
 
 
@@ -16,9 +16,9 @@ class MediaType(BaseModel):
         - https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.3.md#mediaTypeObject
     """
 
-    media_type_schema: Optional[Union[Reference, Schema]] = Field(default=None, alias="schema")
+    media_type_schema: Optional[ReferenceOr[Schema]] = Field(default=None, alias="schema")
     example: Optional[Any] = None
-    examples: Optional[dict[str, Union[Example, Reference]]] = None
+    examples: Optional[dict[str, ReferenceOr[Example]]] = None
     encoding: Optional[dict[str, Encoding]] = None
     model_config = ConfigDict(
         # `Encoding` is not build yet, will rebuild in `__init__.py`:

openapi_python_client/schema/openapi_schema_pydantic/operation.py

@@ -1,11 +1,11 @@
-from typing import Optional, Union
+from typing import Optional
 
 from pydantic import BaseModel, ConfigDict, Field
 
 from .callback import Callback
 from .external_documentation import ExternalDocumentation
 from .parameter import Parameter
-from .reference import Reference
+from .reference import ReferenceOr
 from .request_body import RequestBody
 from .responses import Responses
 from .security_requirement import SecurityRequirement
@@ -25,8 +25,8 @@ class Operation(BaseModel):
     description: Optional[str] = None
     externalDocs: Optional[ExternalDocumentation] = None
     operationId: Optional[str] = None
-    parameters: Optional[list[Union[Parameter, Reference]]] = None
-    request_body: Optional[Union[RequestBody, Reference]] = Field(None, alias="requestBody")
+    parameters: Optional[list[ReferenceOr[Parameter]]] = None
+    request_body: Optional[ReferenceOr[RequestBody]] = Field(None, alias="requestBody")
     responses: Responses
     callbacks: Optional[dict[str, Callback]] = None
 

openapi_python_client/schema/openapi_schema_pydantic/parameter.py

@@ -1,11 +1,11 @@
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 from pydantic import BaseModel, ConfigDict, Field
 
 from ..parameter_location import ParameterLocation
 from .example import Example
 from .media_type import MediaType
-from .reference import Reference
+from .reference import ReferenceOr
 from .schema import Schema
 
 
@@ -30,9 +30,9 @@ class Parameter(BaseModel):
     style: Optional[str] = None
     explode: bool = False
     allowReserved: bool = False
-    param_schema: Optional[Union[Reference, Schema]] = Field(default=None, alias="schema")
+    param_schema: Optional[ReferenceOr[Schema]] = Field(default=None, alias="schema")
     example: Optional[Any] = None
-    examples: Optional[dict[str, Union[Example, Reference]]] = None
+    examples: Optional[dict[str, ReferenceOr[Example]]] = None
     content: Optional[dict[str, MediaType]] = None
     model_config = ConfigDict(
         # `MediaType` is not build yet, will rebuild in `__init__.py`:

openapi_python_client/schema/openapi_schema_pydantic/path_item.py

@@ -1,9 +1,9 @@
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Optional
 
 from pydantic import BaseModel, ConfigDict, Field
 
 from .parameter import Parameter
-from .reference import Reference
+from .reference import ReferenceOr
 from .server import Server
 
 if TYPE_CHECKING:
@@ -34,7 +34,7 @@ class PathItem(BaseModel):
     patch: Optional["Operation"] = None
     trace: Optional["Operation"] = None
     servers: Optional[list[Server]] = None
-    parameters: Optional[list[Union[Parameter, Reference]]] = None
+    parameters: Optional[list[ReferenceOr[Parameter]]] = None
     model_config = ConfigDict(
         # `Operation` is an unresolvable forward reference, will rebuild in `__init__.py`:
         defer_build=True,

openapi_python_client/schema/openapi_schema_pydantic/reference.py

@@ -1,4 +1,7 @@
-from pydantic import BaseModel, ConfigDict, Field
+from typing import Annotated, Any, Literal, TypeVar
+
+from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag
+from pydantic_core import ArgsKwargs
 
 
 class Reference(BaseModel):
@@ -24,3 +27,20 @@ class Reference(BaseModel):
             "examples": [{"$ref": "#/components/schemas/Pet"}, {"$ref": "Pet.json"}, {"$ref": "definitions.json#/Pet"}]
         },
     )
+
+
+T = TypeVar('T')
+
+
+def _reference_discriminator(obj: Any) -> Literal['ref', 'other']:
+    if isinstance(obj, dict):
+        return 'ref' if '$ref' in obj else 'other'
+    if isinstance(obj, ArgsKwargs):
+        return 'ref' if 'ref' in obj.kwargs else 'other'
+    return 'ref' if isinstance(obj, Reference) else 'other'
+
+
+ReferenceOr = Annotated[
+    Annotated[Reference, Tag('ref')] |
+    Annotated[T, Tag('other')],
+    Discriminator(_reference_discriminator)]

openapi_python_client/schema/openapi_schema_pydantic/response.py

@@ -1,11 +1,11 @@
-from typing import Optional, Union
+from typing import Optional
 
 from pydantic import BaseModel, ConfigDict
 
 from .header import Header
 from .link import Link
 from .media_type import MediaType
-from .reference import Reference
+from .reference import ReferenceOr
 
 
 class Response(BaseModel):
@@ -19,9 +19,9 @@ class Response(BaseModel):
     """
 
     description: str
-    headers: Optional[dict[str, Union[Header, Reference]]] = None
+    headers: Optional[dict[str, ReferenceOr[Header]]] = None
     content: Optional[dict[str, MediaType]] = None
-    links: Optional[dict[str, Union[Link, Reference]]] = None
+    links: Optional[dict[str, ReferenceOr[Link]]] = None
     model_config = ConfigDict(
         # `MediaType` is not build yet, will rebuild in `__init__.py`:
         defer_build=True,

openapi_python_client/schema/openapi_schema_pydantic/responses.py

@@ -1,9 +1,8 @@
-from typing import Union
 
-from .reference import Reference
+from .reference import ReferenceOr
 from .response import Response
 
-Responses = dict[str, Union[Response, Reference]]
+Responses = dict[str, ReferenceOr[Response]]
 """
 A container for the expected responses of an operation.
 The container maps a HTTP response code to the expected response.

openapi_python_client/schema/openapi_schema_pydantic/schema.py

@@ -5,7 +5,7 @@
 from ..data_type import DataType
 from .discriminator import Discriminator
 from .external_documentation import ExternalDocumentation
-from .reference import Reference
+from .reference import ReferenceOr
 from .xml import XML
 
 
@@ -38,14 +38,14 @@ class Schema(BaseModel):
     enum: Union[None, list[Any]] = Field(default=None, min_length=1)
     const: Union[None, StrictStr, StrictInt, StrictFloat, StrictBool] = None
     type: Union[DataType, list[DataType], None] = Field(default=None)
-    allOf: list[Union[Reference, "Schema"]] = Field(default_factory=list)
-    oneOf: list[Union[Reference, "Schema"]] = Field(default_factory=list)
-    anyOf: list[Union[Reference, "Schema"]] = Field(default_factory=list)
-    schema_not: Optional[Union[Reference, "Schema"]] = Field(default=None, alias="not")
-    items: Optional[Union[Reference, "Schema"]] = None
-    prefixItems: list[Union[Reference, "Schema"]] = Field(default_factory=list)
-    properties: Optional[dict[str, Union[Reference, "Schema"]]] = None
-    additionalProperties: Optional[Union[bool, Reference, "Schema"]] = None
+    allOf: list[ReferenceOr["Schema"]] = Field(default_factory=list)
+    oneOf: list[ReferenceOr["Schema"]] = Field(default_factory=list)
+    anyOf: list[ReferenceOr["Schema"]] = Field(default_factory=list)
+    schema_not: Optional[ReferenceOr["Schema"]] = Field(default=None, alias="not")
+    items: Optional[ReferenceOr["Schema"]] = None
+    prefixItems: Optional[list[ReferenceOr["Schema"]]] = Field(default_factory=list)
+    properties: Optional[dict[str, ReferenceOr["Schema"]]] = None
+    additionalProperties: Optional[Union[bool, ReferenceOr["Schema"]]] = None
     description: Optional[str] = None
     schema_format: Optional[str] = Field(default=None, alias="format")
     default: Optional[Any] = None

tests/test_schema/test_noisy_refs.py

@@ -0,0 +1,89 @@
+# If a field may be reference (`Union[Reference, OtherType]`) and the dictionary
+# being processed for it contains "$ref", it seems like it should preferentially
+# be parsed as a `Reference`[1].  Since the models are defined with
+# `extra="allow"`, Pydantic won't guarantee this parse if the dictionary is in
+# an unspecified sense a "better match" for `OtherType`[2], e.g., perhaps if it
+# has several more fields matching that type versus the single match for `$ref`.
+#
+# We can use a discriminated union to force parsing these dictionaries as
+# `Reference`s.
+#
+# References:
+#   [1] https://github.com/OAI/OpenAPI-Specification/blob/main/versions/3.0.3.md#reference-object
+#   [2] https://docs.pydantic.dev/latest/concepts/unions/#smart-mode
+import json
+from typing import Annotated, TypeVar, Union, get_args, get_origin
+
+import pytest
+from pydantic import TypeAdapter
+
+from openapi_python_client.schema.openapi_schema_pydantic import (
+    Callback,
+    Example,
+    Header,
+    Link,
+    Parameter,
+    PathItem,
+    Reference,
+    RequestBody,
+    Response,
+    Schema,
+    SecurityScheme,
+)
+
+try:
+    from openapi_python_client.schema.openapi_schema_pydantic.reference import ReferenceOr
+except ImportError:
+    T = TypeVar("T")
+    ReferenceOr = Union[Reference, T]
+
+
+def get_example(base_type):
+    schema = base_type.model_json_schema()
+    print(json.dumps(schema.get("examples", []), indent=4))
+    if "examples" in schema:
+        return schema["examples"][0]
+    if "$defs" in schema:
+        return schema["$defs"][base_type.__name__]["examples"][0]
+    raise TypeError(f"No example found for {base_type.__name__}")
+
+
+def deannotate_type(t):
+    while get_origin(t) is Annotated:
+        t = get_args(t)[0]
+    return t
+
+
+# The following types occur in various models, so we want to make sure they
+# parse properly.  They are verified to /fail/ to parse as of commit 3bd12f86.
+
+@pytest.mark.parametrize(("ref_or_type", "get_example_fn"), [
+    (ReferenceOr[Callback], lambda t: {"test1": get_example(PathItem),
+                                       "test2": get_example(PathItem)}),
+    (ReferenceOr[Example], get_example),
+    (ReferenceOr[Header], get_example),
+    (ReferenceOr[Link], get_example),
+    (ReferenceOr[Parameter], get_example),
+    (ReferenceOr[RequestBody], get_example),
+    (ReferenceOr[Response], get_example),
+    (ReferenceOr[Schema], get_example),
+    (ReferenceOr[SecurityScheme], get_example),
+])
+def test_type(ref_or_type, get_example_fn):
+    base_type = None
+    print(deannotate_type(ref_or_type))
+    for maybe_annotated_type in get_args(deannotate_type(ref_or_type)):
+        each_type = deannotate_type(maybe_annotated_type)
+        if each_type is not Reference:
+            base_type = each_type
+            break
+    assert base_type is not None
+
+    example = get_example_fn(base_type)
+
+    parsed = TypeAdapter(ref_or_type).validate_python(example)
+    assert type(parsed) is get_origin(base_type) or base_type
+
+    example["$ref"] = "ref"
+    parsed = TypeAdapter(ref_or_type).validate_python(example)
+    assert type(parsed) is Reference