✨ Feature: Support pydantic v2 (#45)

Co-authored-by: Ju4tCode <42488585+yanyongyu@users.noreply.github.com>

Author: sudosubin

codegen/__init__.py

@@ -40,7 +40,7 @@ def load_config() -> Config:
     pyproject = tomli.loads(Path("./pyproject.toml").read_text())
     config_dict: Dict[str, Any] = pyproject.get("tool", {}).get("codegen", {})
 
-    return Config.parse_obj(config_dict)
+    return Config.model_validate(config_dict)
 
 
 def build_rest_api(data: OpenAPIData, rest: RestConfig):

codegen/parser/__init__.py

@@ -57,7 +57,7 @@ def parse_openapi_spec(source: Source, rest: RestConfig, config: Config) -> Open
     _st = _schemas.set({})
 
     try:
-        openapi = OpenAPI.parse_obj(source.root)
+        openapi = OpenAPI.model_validate(source.root)
 
         # cache /components/schemas first
         if openapi.components and openapi.components.schemas:

codegen/parser/endpoints/__init__.py

@@ -14,7 +14,7 @@
 
 def parse_endpoint(source: Source, path: str) -> List[EndpointData]:
     data = source.data
-    data = oas.PathItem.parse_obj(data)
+    data = oas.PathItem.model_validate(data)
 
     endpoints: List[EndpointData] = []
 

codegen/parser/endpoints/parameter.py

@@ -1,7 +1,7 @@
 from typing import Union, Literal
 
 import openapi_pydantic as oas
-from pydantic import parse_obj_as
+from pydantic import TypeAdapter
 
 from ...source import Source
 from ..utils import build_prop_name, concat_snake_name
@@ -15,14 +15,14 @@ class Parameter(Property):
 def build_param(source: Source, prefix: str) -> Parameter:
     data = source.data
     try:
-        data = parse_obj_as(Union[oas.Reference, oas.Parameter], data)
+        data = TypeAdapter(Union[oas.Reference, oas.Parameter]).validate_python(data)
     except Exception as e:
         raise TypeError(f"Invalid Parameter from {source.uri}") from e
 
     if isinstance(data, oas.Reference):
         source = source.resolve_ref(data.ref)
         try:
-            data = oas.Parameter.parse_obj(source.data)
+            data = oas.Parameter.model_validate(source.data)
         except Exception as e:
             raise TypeError(f"Invalid Parameter from {source.uri}") from e
 
@@ -38,5 +38,5 @@ def build_param(source: Source, prefix: str) -> Parameter:
         prop_name=build_prop_name(data.name),
         required=data.required,
         schema_data=schema,
-        param_in=data.param_in,  # type: ignore
+        param_in=data.param_in.value,
     )

codegen/parser/endpoints/request_body.py

@@ -1,7 +1,7 @@
 from typing import Set, List, Union, Literal
 
 import openapi_pydantic as oas
-from pydantic import BaseModel, parse_obj_as
+from pydantic import BaseModel, TypeAdapter
 
 from ...source import Source
 from ..utils import concat_snake_name
@@ -58,14 +58,14 @@ def get_using_imports(self) -> Set[str]:
 def build_request_body(source: Source, prefix: str) -> RequestBodyData:
     data = source.data
     try:
-        data = parse_obj_as(Union[oas.Reference, oas.RequestBody], data)
+        data = TypeAdapter(Union[oas.Reference, oas.RequestBody]).validate_python(data)
     except Exception as e:
         raise TypeError(f"Invalid RequestBody from {source.uri}") from e
 
     if isinstance(data, oas.Reference):
         source = source.resolve_ref(data.ref)
         try:
-            data = oas.RequestBody.parse_obj(source.data)
+            data = oas.RequestBody.model_validate(source.data)
         except Exception as e:
             raise TypeError(f"Invalid RequestBody from {source.uri}") from e
 

codegen/parser/endpoints/response.py

@@ -1,7 +1,7 @@
 from typing import Set, Union, Optional
 
 import openapi_pydantic as oas
-from pydantic import BaseModel, parse_obj_as
+from pydantic import BaseModel, TypeAdapter
 
 from ...source import Source
 from ..schemas import SchemaData, parse_schema
@@ -20,14 +20,14 @@ def get_using_imports(self) -> Set[str]:
 def build_response(source: Source, prefix: str) -> ResponseData:
     data = source.data
     try:
-        data = parse_obj_as(Union[oas.Reference, oas.Response], data)
+        data = TypeAdapter(Union[oas.Reference, oas.Response]).validate_python(data)
     except Exception as e:
         raise TypeError(f"Invalid Response from {source.uri}") from e
 
     if isinstance(data, oas.Reference):
         source = source.resolve_ref(data.ref)
         try:
-            data = oas.Response.parse_obj(source.data)
+            data = oas.Response.model_validate(source.data)
         except Exception as e:
             raise TypeError(f"Invalid Response from {source.uri}") from e
 

codegen/parser/schemas/__init__.py

@@ -1,7 +1,7 @@
 from typing import Union, Optional
 
 import openapi_pydantic as oas
-from pydantic import parse_obj_as
+from pydantic import TypeAdapter
 
 from ...source import Source
 from ..utils import schema_from_source
@@ -21,14 +21,15 @@
 from .schema import StringSchema as StringSchema
 from .. import add_schema, get_schema, get_schemas
 from .schema import DateTimeSchema as DateTimeSchema
+from .schema import UniqueListSchema as UniqueListSchema
 
 
 def parse_schema(
     source: Source, class_name: str, base_source: Optional[Source] = None
 ) -> SchemaData:
     data = source.data
     try:
-        data = parse_obj_as(Union[oas.Reference, oas.Schema], data)
+        data = TypeAdapter(Union[oas.Reference, oas.Schema]).validate_python(data)
     except Exception as e:
         raise TypeError(f"Invalid Schema from {source.uri}") from e
 

codegen/parser/schemas/any_schema.py

@@ -6,7 +6,7 @@
 
 def build_any_schema(source: Source) -> AnySchema:
     try:
-        data = oas.Schema.parse_obj(source.data)
+        data = oas.Schema.model_validate(source.data)
     except Exception as e:
         raise TypeError(f"Invalid Schema from {source.uri}") from e
 

codegen/parser/schemas/bool_schema.py

@@ -7,7 +7,7 @@
 
 def build_bool_schema(source: Source) -> BoolSchema:
     try:
-        data = oas.Schema.parse_obj(source.data)
+        data = oas.Schema.model_validate(source.data)
     except Exception as e:
         raise TypeError(f"Invalid Schema from {source.uri}") from e
 

codegen/parser/schemas/list_schema.py

@@ -1,26 +1,30 @@
-from typing import Optional
+from typing import Union, Optional
 
 from . import parse_schema
 from ...source import Source
-from .schema import ListSchema
+from .schema import ListSchema, UniqueListSchema
 from ..utils import concat_snake_name, schema_from_source
 
 
 def build_list_schema(
     source: Source, class_name: str, base_source: Optional[Source] = None
-) -> ListSchema:
+) -> Union[ListSchema, UniqueListSchema]:
     data = schema_from_source(source)
     base_schema = schema_from_source(base_source) if base_source else None
 
-    return ListSchema(
+    if data.uniqueItems or (base_schema and base_schema.uniqueItems):
+        schema_class = UniqueListSchema
+    else:
+        schema_class = ListSchema
+
+    return schema_class(
         title=data.title,
         description=data.description,
         default=data.default,
         examples=data.examples or (data.example and [data.example]),
         item_schema=parse_schema(
             source / "items", concat_snake_name(class_name, "items")
         ),
-        min_items=data.minItems or (base_schema and base_schema.minItems),
-        max_items=data.maxItems or (base_schema and base_schema.maxItems),
-        unique_items=data.uniqueItems or (base_schema and base_schema.uniqueItems),
+        min_length=data.minItems or (base_schema and base_schema.minItems),
+        max_length=data.maxItems or (base_schema and base_schema.maxItems),
     )

codegen/parser/schemas/model_schema.py

@@ -26,6 +26,7 @@
     UnionSchema,
     StringSchema,
     DateTimeSchema,
+    UniqueListSchema,
 )
 
 ST = TypeVar("ST", bound=SchemaData)
@@ -61,7 +62,7 @@ def _is_union_subset(first: SchemaData, second: SchemaData) -> Optional[SchemaDa
         second_schemas if len(first_schemas) <= len(second_schemas) else first_schemas
     )
     for schema in one_schemas:
-        if schema not in another_schemas:
+        if schema.model_dump() not in [s.model_dump() for s in another_schemas]:
             return
     return first if len(first_schemas) <= len(second_schemas) else second
 
@@ -157,6 +158,27 @@ def _is_list_merge(
         )
 
 
+def _is_unique_list_merge(
+    source: Source, name: str, prefix: str, first: SchemaData, second: SchemaData
+) -> Optional[UniqueListSchema]:
+    if (
+        (isinstance(first, UniqueListSchema) and isinstance(second, ListSchema))
+        or (isinstance(first, ListSchema) and isinstance(second, UniqueListSchema))
+        or (
+            isinstance(first, UniqueListSchema) and isinstance(second, UniqueListSchema)
+        )
+    ):
+        return UniqueListSchema(
+            title=first.title,
+            description=first.description,
+            default=first.default,
+            examples=first.examples,
+            item_schema=_merge_schema(
+                source, name, prefix, first.item_schema, second.item_schema
+            ),
+        )
+
+
 def _merge_schema(
     source: Source, name: str, prefix: str, first: SchemaData, second: SchemaData
 ):
@@ -167,6 +189,7 @@ def _merge_schema(
         or _is_enum_subset(first, second)
         or _is_model_merge(source, name, prefix, first, second)
         or _is_list_merge(source, name, prefix, first, second)
+        or _is_unique_list_merge(source, name, prefix, first, second)
     ):
         return schema
     raise RuntimeError(f"Cannot merge schema for {name}: {first!r}; {second!r}")

codegen/parser/schemas/none_schema.py

@@ -6,7 +6,7 @@
 
 def build_none_schema(source: Source) -> NoneSchema:
     try:
-        data = oas.Schema.parse_obj(source.data)
+        data = oas.Schema.model_validate(source.data)
     except Exception as e:
         raise TypeError(f"Invalid Schema from {source.uri}") from e