Satisfy mypy with either version of Pydantic

Author: hathawsh

openapi_pydantic/__init__.py

@@ -33,5 +33,6 @@
 from .v3 import ServerVariable as ServerVariable
 from .v3 import Tag as Tag
 from .v3 import parse_obj as parse_obj
+from .v3 import schema_validate as schema_validate
 
 logging.getLogger(__name__).addHandler(logging.NullHandler())

openapi_pydantic/compat.py

@@ -1,5 +1,7 @@
 """Compatibility layer to make this package usable with Pydantic 1 or 2"""
 
+from typing import TYPE_CHECKING
+
 from pydantic.version import VERSION as PYDANTIC_VERSION
 
 PYDANTIC_MAJOR_VERSION = int(PYDANTIC_VERSION.split(".", 1)[0])
@@ -9,45 +11,107 @@
 else:
     PYDANTIC_V2 = False
 
-if PYDANTIC_V2:
-    from typing import Literal
+if TYPE_CHECKING:
+    # Provide stubs for either version of Pydantic
+
+    from enum import Enum
+    from typing import Any, Literal, Type, TypedDict
+
+    from pydantic import BaseModel
+    from pydantic import ConfigDict as PydanticConfigDict
+
+    def ConfigDict(
+        extra: Literal["allow", "ignore", "forbid"] = "allow",
+        json_schema_extra: dict[str, Any] | None = None,
+        populate_by_name: bool = True,
+    ) -> PydanticConfigDict:
+        ...
+
+    class Extra(Enum):
+        allow = "allow"
+        ignore = "ignore"
+        forbid = "forbid"
+
+    class RootModel(BaseModel):
+        ...
+
+    JsonSchemaMode = Literal["validation", "serialization"]
+
+    def models_json_schema(
+        models: list[tuple[Type[BaseModel], JsonSchemaMode]],
+        *,
+        by_alias: bool = True,
+        ref_template: str = "#/$defs/{model}",
+    ) -> tuple[dict, dict[str, Any]]:
+        ...
+
+    def v1_schema(
+        models: list[Type[BaseModel]],
+        *,
+        by_alias: bool = True,
+        ref_prefix: str = "#/$defs",
+    ) -> dict[str, Any]:
+        ...
+
+    DEFS_KEY = "$defs"
 
-    from pydantic import ConfigDict
-    from pydantic.json_schema import JsonSchemaMode, models_json_schema  # type: ignore
+    class MinLengthArg(TypedDict):
+        pass
+
+    def min_length_arg(min_length: int) -> MinLengthArg:
+        ...
+
+elif PYDANTIC_V2:
+    from typing import Literal, TypedDict
+
+    from pydantic import ConfigDict, RootModel
+    from pydantic.json_schema import JsonSchemaMode, models_json_schema
 
     # Pydantic 2 renders JSON schemas using the keyword "$defs"
     DEFS_KEY = "$defs"
 
-    # Add V1 stubs to this module, but hide them from typing
-    globals().update(
-        {
-            "Extra": None,
-            "v1_schema": None,
-        }
-    )
+    class MinLengthArg(TypedDict):
+        min_length: int
+
+    def min_length_arg(min_length: int) -> MinLengthArg:
+        return {"min_length": min_length}
+
+    # Create V1 stubs
+    Extra = None
+    v1_schema = None
+
 
 else:
+    from typing import TypedDict
+
     from pydantic import Extra
     from pydantic.schema import schema as v1_schema
 
     # Pydantic 1 renders JSON schemas using the keyword "definitions"
     DEFS_KEY = "definitions"
 
-    # Add V2 stubs to this module, but hide them from typing
-    globals().update(
-        {
-            "ConfigDict": None,
-            "Literal": None,
-            "models_json_schema": None,
-            "JsonSchemaMode": None,
-        }
-    )
+    class MinLengthArg(TypedDict):
+        min_items: int
+
+    def min_length_arg(min_length: int) -> MinLengthArg:
+        return {"min_items": min_length}
+
+    # Create V2 stubs
+    ConfigDict = None
+    Literal = None
+    models_json_schema = None
+    JsonSchemaMode = None
+    RootModel = None
+
 
 __all__ = [
     "Literal",
     "ConfigDict",
     "JsonSchemaMode",
     "models_json_schema",
+    "RootModel",
     "Extra",
     "v1_schema",
+    "DEFS_KEY",
+    "min_length_arg",
 ]

openapi_pydantic/util.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Generic, List, Optional, Set, Type, TypeVar
+from typing import Any, Generic, List, Optional, Set, Type, TypeVar, cast
 
 from pydantic import BaseModel
 
@@ -11,7 +11,7 @@
     v1_schema,
 )
 
-from . import Components, OpenAPI, Reference, Schema
+from . import Components, OpenAPI, Reference, Schema, schema_validate
 
 logger = logging.getLogger(__name__)
 
@@ -32,13 +32,16 @@ def get_mode(
 ) -> JsonSchemaMode:
     """Get the JSON schema mode for a model class.
 
-    The mode can be either "serialization" or "validation". In validation mode,
+    The mode can be either "validation" or "serialization". In validation mode,
     computed fields are dropped and optional fields remain optional. In
     serialization mode, computed and optional fields are required.
     """
     if not hasattr(cls, "model_config"):
         return default
-    return cls.model_config.get("json_schema_mode", default)
+    mode = cls.model_config.get("json_schema_mode", default)
+    if mode not in ("validation", "serialization"):
+        raise ValueError(f"invalid json_schema_mode: {mode}")
+    return cast(JsonSchemaMode, mode)
 
 
 def construct_open_api_with_schema_class(
@@ -62,10 +65,8 @@ def construct_open_api_with_schema_class(
              If there is no update in "#/components/schemas" values, the original
              `open_api` will be returned.
     """
-    if PYDANTIC_V2:
-        new_open_api = open_api.model_copy(deep=True)
-    else:
-        new_open_api = open_api.copy(deep=True)
+    copy_func = getattr(open_api, "model_copy" if PYDANTIC_V2 else "copy")
+    new_open_api: OpenAPI = copy_func(deep=True)
 
     if scan_for_pydantic_schema_reference:
         extracted_schema_classes = _handle_pydantic_schema(new_open_api)
@@ -80,7 +81,7 @@ def construct_open_api_with_schema_class(
         return open_api
 
     schema_classes.sort(key=lambda x: x.__name__)
-    logger.debug(f"schema_classes{schema_classes}")
+    logger.debug("schema_classes: %s", schema_classes)
 
     # update new_open_api with new #/components/schemas
     if PYDANTIC_V2:
@@ -94,7 +95,6 @@ def construct_open_api_with_schema_class(
             schema_classes, by_alias=by_alias, ref_prefix=ref_prefix
         )
 
-    schema_validate = Schema.model_validate if PYDANTIC_V2 else Schema.parse_obj
     if not new_open_api.components:
         new_open_api.components = Components()
     if new_open_api.components.schemas:
@@ -111,6 +111,8 @@ def construct_open_api_with_schema_class(
             }
         )
     else:
+        for key, schema_dict in schema_definitions[DEFS_KEY].items():
+            schema_validate(schema_dict)
         new_open_api.components.schemas = {
             key: schema_validate(schema_dict)
             for key, schema_dict in schema_definitions[DEFS_KEY].items()
@@ -136,13 +138,13 @@ def _handle_pydantic_schema(open_api: OpenAPI) -> List[Type[BaseModel]]:
 
     def _traverse(obj: Any) -> None:
         if isinstance(obj, BaseModel):
-            fields = obj.model_fields_set if PYDANTIC_V2 else obj.__fields_set__
+            fields = getattr(
+                obj, "model_fields_set" if PYDANTIC_V2 else "__fields_set__"
+            )
             for field in fields:
                 child_obj = obj.__getattribute__(field)
                 if isinstance(child_obj, PydanticSchema):
-                    logger.debug(
-                        f"PydanticSchema found in {obj.__repr_name__()}: {child_obj}"
-                    )
+                    logger.debug("PydanticSchema found in %s: %s", obj, child_obj)
                     obj.__setattr__(field, _construct_ref_obj(child_obj))
                     pydantic_types.add(child_obj.schema_class)
                 else:
@@ -169,6 +171,6 @@ def _traverse(obj: Any) -> None:
 
 
 def _construct_ref_obj(pydantic_schema: PydanticSchema[PydanticType]) -> Reference:
-    ref_obj = Reference(ref=ref_prefix + pydantic_schema.schema_class.__name__)
+    ref_obj = Reference(**{"$ref": ref_prefix + pydantic_schema.schema_class.__name__})
     logger.debug(f"ref_obj={ref_obj}")
     return ref_obj

openapi_pydantic/v3/__init__.py

@@ -31,3 +31,4 @@
 from .v3_1_0 import Server as Server
 from .v3_1_0 import ServerVariable as ServerVariable
 from .v3_1_0 import Tag as Tag
+from .v3_1_0 import schema_validate as schema_validate

openapi_pydantic/v3/parser.py

@@ -1,4 +1,4 @@
-from typing import Any, Union
+from typing import TYPE_CHECKING, Any, Union
 
 from pydantic import BaseModel, Field
 
@@ -9,21 +9,25 @@
 
 OpenAPIv3 = Union[OpenAPIv3_1, OpenAPIv3_0]
 
-if PYDANTIC_V2:
+if TYPE_CHECKING:
+
+    def parse_obj(data: Any) -> OpenAPIv3:
+        """Parse a raw object into an OpenAPI model with version inference."""
+        ...
+
+elif PYDANTIC_V2:
     from pydantic import RootModel
 
-    class _OpenAPIV2(RootModel):
+    class _OpenAPI(RootModel):
         root: OpenAPIv3 = Field(discriminator="openapi")
 
+    def parse_obj(data: Any) -> OpenAPIv3:
+        return _OpenAPI.model_validate(data).root
+
 else:
 
-    class _OpenAPIV1(BaseModel):
+    class _OpenAPI(BaseModel):
         __root__: OpenAPIv3 = Field(discriminator="openapi")
 
-
-def parse_obj(data: Any) -> OpenAPIv3:
-    """Parse a raw object into an OpenAPI model with version inference."""
-    if PYDANTIC_V2:
-        return _OpenAPIV2.model_validate(data).root
-    else:
-        return _OpenAPIV1.parse_obj(data).__root__
+    def parse_obj(data: Any) -> OpenAPIv3:
+        return _OpenAPI.parse_obj(data).__root__

openapi_pydantic/v3/v3_0_3/__init__.py

@@ -6,6 +6,8 @@
 https://github.com/OAI/OpenAPI-Specification/blob/master/versions/3.0.3.md#table-of-contents
 """
 
+from typing import TYPE_CHECKING
+
 from openapi_pydantic.compat import PYDANTIC_V2
 
 from .callback import Callback as Callback
@@ -34,20 +36,24 @@
 from .response import Response as Response
 from .responses import Responses as Responses
 from .schema import Schema as Schema
+from .schema import schema_validate as schema_validate
 from .security_requirement import SecurityRequirement as SecurityRequirement
 from .security_scheme import SecurityScheme as SecurityScheme
 from .server import Server as Server
 from .server_variable import ServerVariable as ServerVariable
 from .tag import Tag as Tag
 from .xml import XML as XML
 
-# resolve forward references
-if PYDANTIC_V2:
+if TYPE_CHECKING:
+    pass
+elif PYDANTIC_V2:
+    # resolve forward references
     Encoding.model_rebuild()
     OpenAPI.model_rebuild()
     Components.model_rebuild()
     Operation.model_rebuild()
 else:
+    # resolve forward references
     Encoding.update_forward_refs(Header=Header)
     Schema.update_forward_refs()
     Operation.update_forward_refs(PathItem=PathItem)

openapi_pydantic/v3/v3_0_3/header.py

@@ -1,3 +1,5 @@
+from typing import TYPE_CHECKING
+
 from pydantic import Field
 
 from openapi_pydantic.compat import PYDANTIC_V2, ConfigDict, Extra, Literal
@@ -12,10 +14,9 @@
 ]
 
 
-if PYDANTIC_V2:
-    LiteralEmptyString = Literal[""]
+if TYPE_CHECKING:
 
-    class HeaderV2(Parameter):
+    class Header(Parameter):
         """
         The Header Object follows the structure of the
         [Parameter Object](#parameterObject) with the following changes:
@@ -27,6 +28,15 @@ class HeaderV2(Parameter):
            to a location of `header` (for example, [`style`](#parameterStyle)).
         """
 
+        name: str = Field(default="")
+        param_in: ParameterLocation = Field(
+            default=ParameterLocation.HEADER, alias="in"
+        )
+
+elif PYDANTIC_V2:
+    LiteralEmptyString = Literal[""]
+
+    class Header(Parameter):
         name: LiteralEmptyString = Field(default="")
         param_in: Literal[ParameterLocation.HEADER] = Field(
             default=ParameterLocation.HEADER, alias="in"
@@ -38,23 +48,9 @@ class HeaderV2(Parameter):
             json_schema_extra={"examples": _examples},
         )
 
-    Header = HeaderV2
-    Header.__name__ = "Header"
-
 else:
 
-    class HeaderV1(Parameter):
-        """
-        The Header Object follows the structure of the
-        [Parameter Object](#parameterObject) with the following changes:
-
-        1. `name` MUST NOT be specified, it is given in the corresponding
-           `headers` map.
-        2. `in` MUST NOT be specified, it is implicitly in `header`.
-        3. All traits that are affected by the location MUST be applicable
-           to a location of `header` (for example, [`style`](#parameterStyle)).
-        """
-
+    class Header(Parameter):
         name: str = Field(default="", const=True)
         param_in: ParameterLocation = Field(
             default=ParameterLocation.HEADER, const=True, alias="in"
@@ -64,6 +60,3 @@ class Config:
             extra = Extra.allow
             allow_population_by_field_name = True
             schema_extra = {"examples": _examples}
-
-    Header = HeaderV1
-    Header.__name__ = "Header"