Merge branch 'main' into required-nullable-annotated

Author: koxudaxi

docs/root-model-and-type-alias.md

@@ -10,12 +10,23 @@ When a schema defines a simple type (not an object with properties), `datamodel-
 - A RootModel or type alias is also generated for the main schema, allowing you to define a single type alias from a schema file (e.g. `model.json` containing `{"title": "MyString", "type": "string"}`)
 - Type aliases cannot be combined with `Annotated` for Pydantic v1
 
-## Pydantic v1 vs v2
+## Type Alias Behavior by Output Type and Python Version
 
-The type of type alias generated depends on the Pydantic version:
+The type of type alias generated depends on the output model type and target Python version:
 
-- **Pydantic v2**: Uses `TypeAliasType` (Python 3.9-3.11) or `type` statement (Python 3.12+)
-- **Pydantic v1**: Uses `TypeAlias` because Pydantic v1 cannot handle `TypeAliasType` objects
+| Output Type | Python 3.12+ | Python 3.10-3.11 | Python 3.9 |
+|-------------|--------------|------------------|------------|
+| **Pydantic v2** | `type` statement | `TypeAliasType` (typing_extensions) | `TypeAliasType` (typing_extensions) |
+| **Pydantic v1** | `TypeAlias` | `TypeAlias` | `TypeAlias` (typing_extensions) |
+| **TypedDict** | `type` statement | `TypeAlias` | `TypeAlias` (typing_extensions) |
+| **dataclasses** | `type` statement | `TypeAlias` | `TypeAlias` (typing_extensions) |
+| **msgspec** | `type` statement | `TypeAlias` | `TypeAlias` (typing_extensions) |
+
+**Why the difference?**
+
+- **Pydantic v2** requires `TypeAliasType` because it cannot properly handle `TypeAlias` annotations
+- **Other output types** (TypedDict, dataclasses, msgspec) use `TypeAlias` for better compatibility with libraries that may not expect `TypeAliasType` objects
+- **Python 3.12+** uses the native `type` statement for all output types
 
 ## Example
 

src/datamodel_code_generator/model/__init__.py

@@ -55,28 +55,28 @@ def get_data_model_types(
     )
     from .types import DataTypeManager  # noqa: PLC0415
 
-    # Pydantic v1 does not support TypeAliasType type, fallback to TypeAlias
-    if data_model_type == DataModelType.PydanticBaseModel:
-        if target_python_version.has_type_alias:
-            # Python 3.10+: typing.TypeAlias
-            type_alias_class = type_alias.TypeAlias
-            scalar_class = scalar.DataTypeScalar
-            union_class = union.DataTypeUnion
+    # Pydantic v2 requires TypeAliasType; other output types use TypeAlias for better compatibility
+    if data_model_type == DataModelType.PydanticV2BaseModel:
+        if target_python_version.has_type_statement:
+            type_alias_class = type_alias.TypeStatement
+            scalar_class = scalar.DataTypeScalarTypeStatement
+            union_class = union.DataTypeUnionTypeStatement
         else:
-            # Python 3.9: typing_extensions.TypeAlias
-            type_alias_class = type_alias.TypeAliasBackport
-            scalar_class = scalar.DataTypeScalarBackport
-            union_class = union.DataTypeUnionBackport
+            type_alias_class = type_alias.TypeAliasTypeBackport
+            scalar_class = scalar.DataTypeScalarTypeBackport
+            union_class = union.DataTypeUnionTypeBackport
     elif target_python_version.has_type_statement:
-        # Python 3.12+ with Pydantic v2 or other formats: Use type statement
         type_alias_class = type_alias.TypeStatement
         scalar_class = scalar.DataTypeScalarTypeStatement
         union_class = union.DataTypeUnionTypeStatement
+    elif target_python_version.has_type_alias:
+        type_alias_class = type_alias.TypeAlias
+        scalar_class = scalar.DataTypeScalar
+        union_class = union.DataTypeUnion
     else:
-        # Python 3.9-3.11 with Pydantic v2 or other formats: Use TypeAliasType
-        type_alias_class = type_alias.TypeAliasTypeBackport
-        scalar_class = scalar.DataTypeScalarTypeBackport
-        union_class = union.DataTypeUnionTypeBackport
+        type_alias_class = type_alias.TypeAliasBackport
+        scalar_class = scalar.DataTypeScalarBackport
+        union_class = union.DataTypeUnionBackport
 
     if data_model_type == DataModelType.PydanticBaseModel:
         return DataModelSet(

src/datamodel_code_generator/model/base.py

@@ -340,6 +340,7 @@ class DataModel(TemplateBase, Nullable, ABC):
     TEMPLATE_FILE_PATH: ClassVar[str] = ""
     BASE_CLASS: ClassVar[str] = ""
     DEFAULT_IMPORTS: ClassVar[tuple[Import, ...]] = ()
+    IS_ALIAS: bool = False
 
     def __init__(  # noqa: PLR0913
         self,
@@ -523,6 +524,11 @@ def all_data_types(self) -> Iterator[DataType]:
             yield from field.data_type.all_data_types
         yield from self.base_classes
 
+    @property
+    def is_alias(self) -> bool:
+        """Whether is a type alias (i.e. not an instance of BaseModel/RootModel)."""
+        return self.IS_ALIAS
+
     @property
     def nullable(self) -> bool:
         """Check if this model is nullable."""

src/datamodel_code_generator/model/type_alias.py

@@ -22,6 +22,8 @@
 class TypeAliasBase(DataModel):
     """Base class for all type alias implementations."""
 
+    IS_ALIAS: bool = True
+
     @property
     def imports(self) -> tuple[Import, ...]:
         """Get imports including Annotated if needed."""

src/datamodel_code_generator/parser/base.py

@@ -132,6 +132,28 @@ def dump_templates(templates: list[DataModel]) -> str:
 MAX_RECURSION_COUNT: int = sys.getrecursionlimit()
 
 
+def add_model_path_to_list(
+    paths: list[str] | None,
+    model: DataModel,
+    /,
+) -> list[str]:
+    """
+    Auxiliary method which adds model path to list, provided the following hold.
+
+    - model is not a type alias
+    - path is not already in the list.
+
+    """
+    if paths is None:
+        paths = []
+    if model.is_alias:
+        return paths
+    if (path := model.path) in paths:
+        return paths
+    paths.append(path)
+    return paths
+
+
 def sort_data_models(  # noqa: PLR0912, PLR0915
     unsorted_data_models: list[DataModel],
     sorted_data_models: SortedDataModels | None = None,
@@ -151,13 +173,13 @@ def sort_data_models(  # noqa: PLR0912, PLR0915
             sorted_data_models[model.path] = model
         elif model.path in model.reference_classes and len(model.reference_classes) == 1:  # only self-referencing
             sorted_data_models[model.path] = model
-            require_update_action_models.append(model.path)
+            add_model_path_to_list(require_update_action_models, model)
         elif (
             not model.reference_classes - {model.path} - set(sorted_data_models)
         ):  # reference classes have been resolved
             sorted_data_models[model.path] = model
             if model.path in model.reference_classes:
-                require_update_action_models.append(model.path)
+                add_model_path_to_list(require_update_action_models, model)
         else:
             unresolved_references.append(model)
     if unresolved_references:
@@ -206,11 +228,11 @@ def sort_data_models(  # noqa: PLR0912, PLR0915
             if not unresolved_model:
                 sorted_data_models[model.path] = model
                 if update_action_parent:
-                    require_update_action_models.append(model.path)
+                    add_model_path_to_list(require_update_action_models, model)
                 continue
             if not unresolved_model - unsorted_data_model_names:
                 sorted_data_models[model.path] = model
-                require_update_action_models.append(model.path)
+                add_model_path_to_list(require_update_action_models, model)
                 continue
             # unresolved
             unresolved_classes = ", ".join(
@@ -1003,7 +1025,7 @@ def __reuse_model(self, models: list[DataModel], require_update_action_models: l
                         custom_template_dir=model._custom_template_dir,  # noqa: SLF001
                     )
                     if cached_model_reference.path in require_update_action_models:
-                        require_update_action_models.append(inherited_model.path)
+                        add_model_path_to_list(require_update_action_models, inherited_model)
                     models.insert(index, inherited_model)
                     models.remove(model)
 
@@ -1341,6 +1363,39 @@ def __alias_shadowed_imports(  # noqa: PLR6301
                             reference_path=model_field.data_type.import_.reference_path,
                         )
 
+    @classmethod
+    def _collect_used_names_from_models(cls, models: list[DataModel]) -> set[str]:
+        """Collect identifiers referenced by models before rendering."""
+        names: set[str] = set()
+
+        def add(name: str | None) -> None:
+            if not name:
+                return
+            # first segment is sufficient to match import target or alias
+            names.add(name.split(".")[0])
+
+        def walk_data_type(data_type: DataType) -> None:
+            add(data_type.alias or data_type.type)
+            if data_type.reference:
+                add(data_type.reference.short_name)
+            for child in data_type.data_types:
+                walk_data_type(child)
+            if data_type.dict_key:
+                walk_data_type(data_type.dict_key)
+
+        for model in models:
+            add(model.class_name)
+            add(model.duplicate_class_name)
+            for base in model.base_classes:
+                add(base.type_hint)
+            for import_ in model.imports:
+                add(import_.alias or import_.import_.split(".")[-1])
+            for field in model.fields:
+                add(field.name)
+                add(field.alias)
+                walk_data_type(field.data_type)
+        return names
+
     def parse(  # noqa: PLR0912, PLR0914, PLR0915
         self,
         with_import: bool | None = True,  # noqa: FBT001, FBT002
@@ -1466,12 +1521,14 @@ class Processed(NamedTuple):
 
         for processed_model in processed_models:
             # postprocess imports to remove unused imports.
-            model_code = str("\n".join([str(m) for m in processed_model.models]))
+            used_names = self._collect_used_names_from_models(processed_model.models)
             unused_imports = [
                 (from_, import_)
                 for from_, imports_ in processed_model.imports.items()
                 for import_ in imports_
-                if import_ not in model_code
+                if not {processed_model.imports.alias.get(from_, {}).get(import_, import_), import_}.intersection(
+                    used_names
+                )
             ]
             for from_, import_ in unused_imports:
                 processed_model.imports.remove(Import(from_=from_, import_=import_))

src/datamodel_code_generator/parser/openapi.py

@@ -344,6 +344,8 @@ def __init__(  # noqa: PLR0913
         )
         self.open_api_scopes: list[OpenAPIScope] = openapi_scopes or [OpenAPIScope.Schemas]
         self.include_path_parameters: bool = include_path_parameters
+        self._discriminator_schemas: dict[str, dict[str, Any]] = {}
+        self._discriminator_subtypes: dict[str, list[str]] = defaultdict(list)
 
     def get_ref_model(self, ref: str) -> dict[str, Any]:
         """Resolve a reference to its model definition."""
@@ -362,6 +364,49 @@ def get_data_type(self, obj: JsonSchemaObject) -> DataType:
 
         return super().get_data_type(obj)
 
+    def _get_discriminator_union_type(self, ref: str) -> DataType | None:
+        """Create a union type for discriminator subtypes if available."""
+        subtypes = self._discriminator_subtypes.get(ref, [])
+        if not subtypes:
+            return None
+        refs = map(self.model_resolver.add_ref, subtypes)
+        return self.data_type(data_types=[self.data_type(reference=r) for r in refs])
+
+    def get_ref_data_type(self, ref: str) -> DataType:
+        """Get data type for a reference, handling discriminator polymorphism."""
+        if ref in self._discriminator_schemas and (union_type := self._get_discriminator_union_type(ref)):
+            return union_type
+        return super().get_ref_data_type(ref)
+
+    def parse_object_fields(
+        self,
+        obj: JsonSchemaObject,
+        path: list[str],
+        module_name: Optional[str] = None,  # noqa: UP045
+    ) -> list[DataModelFieldBase]:
+        """Parse object fields, adding discriminator info for allOf polymorphism."""
+        fields = super().parse_object_fields(obj, path, module_name)
+        properties = obj.properties or {}
+
+        result_fields: list[DataModelFieldBase] = []
+        for field_obj in fields:
+            field = properties.get(field_obj.original_name)
+
+            if (
+                isinstance(field, JsonSchemaObject)
+                and field.ref
+                and (discriminator := self._discriminator_schemas.get(field.ref))
+            ):
+                new_field_type = self._get_discriminator_union_type(field.ref) or field_obj.data_type
+                field_obj = self.data_model_field_type(**{  # noqa: PLW2901
+                    **field_obj.__dict__,
+                    "data_type": new_field_type,
+                    "extras": {**field_obj.extras, "discriminator": discriminator},
+                })
+            result_fields.append(field_obj)
+
+        return result_fields
+
     def resolve_object(self, obj: ReferenceObject | BaseModelT, object_type: type[BaseModelT]) -> BaseModelT:
         """Resolve a reference object to its actual type or return the object as-is."""
         if isinstance(obj, ReferenceObject):
@@ -651,6 +696,7 @@ def parse_raw(self) -> None:  # noqa: PLR0912, PLR0915
 
             specification: dict[str, Any] = load_yaml_dict(source.text)
             self.raw_obj = specification
+            self._collect_discriminator_schemas()
             schemas: dict[str, Any] = specification.get("components", {}).get("schemas", {})
             security: list[dict[str, list[str]]] | None = specification.get("security")
             if OpenAPIScope.Schemas in self.open_api_scopes:
@@ -727,3 +773,25 @@ def parse_raw(self) -> None:  # noqa: PLR0912, PLR0915
                         )
 
         self._resolve_unparsed_json_pointer()
+
+    def _collect_discriminator_schemas(self) -> None:
+        """Collect schemas with discriminators but no oneOf/anyOf, and find their subtypes."""
+        schemas: dict[str, Any] = self.raw_obj.get("components", {}).get("schemas", {})
+
+        for schema_name, schema in schemas.items():
+            discriminator = schema.get("discriminator")
+            if not discriminator:
+                continue
+
+            if schema.get("oneOf") or schema.get("anyOf"):
+                continue
+
+            ref = f"#/components/schemas/{schema_name}"
+            self._discriminator_schemas[ref] = discriminator
+
+        for schema_name, schema in schemas.items():
+            for all_of_item in schema.get("allOf", []):
+                ref_in_allof = all_of_item.get("$ref")
+                if ref_in_allof and ref_in_allof in self._discriminator_schemas:
+                    subtype_ref = f"#/components/schemas/{schema_name}"
+                    self._discriminator_subtypes[ref_in_allof].append(subtype_ref)

tests/data/expected/main/graphql/simple_star_wars_dataclass.py

@@ -7,27 +7,27 @@
 from dataclasses import dataclass
 from typing import List, Literal, Optional
 
-from typing_extensions import TypeAliasType
+from typing_extensions import TypeAlias
 
-Boolean = TypeAliasType("Boolean", bool)
+Boolean: TypeAlias = bool
 """
 The `Boolean` scalar type represents `true` or `false`.
 """
 
 
-ID = TypeAliasType("ID", str)
+ID: TypeAlias = str
 """
 The `ID` scalar type represents a unique identifier, often used to refetch an object or as key for a cache. The ID type appears in a JSON response as a String; however, it is not intended to be human-readable. When expected as an input type, any string (such as `"4"`) or integer (such as `4`) input value will be accepted as an ID.
 """
 
 
-Int = TypeAliasType("Int", int)
+Int: TypeAlias = int
 """
 The `Int` scalar type represents non-fractional signed whole numeric values. Int can represent values between -(2^31) and 2^31 - 1.
 """
 
 
-String = TypeAliasType("String", str)
+String: TypeAlias = str
 """
 The `String` scalar type represents textual data, represented as UTF-8 character sequences. The String type is most often used by GraphQL to represent free-form human-readable text.
 """

tests/data/expected/main/graphql/simple_star_wars_dataclass_arguments.py

@@ -7,27 +7,27 @@
 from dataclasses import dataclass
 from typing import List, Literal, Optional
 
-from typing_extensions import TypeAliasType
+from typing_extensions import TypeAlias
 
-Boolean = TypeAliasType("Boolean", bool)
+Boolean: TypeAlias = bool
 """
 The `Boolean` scalar type represents `true` or `false`.
 """
 
 
-ID = TypeAliasType("ID", str)
+ID: TypeAlias = str
 """
 The `ID` scalar type represents a unique identifier, often used to refetch an object or as key for a cache. The ID type appears in a JSON response as a String; however, it is not intended to be human-readable. When expected as an input type, any string (such as `"4"`) or integer (such as `4`) input value will be accepted as an ID.
 """
 
 
-Int = TypeAliasType("Int", int)
+Int: TypeAlias = int
 """
 The `Int` scalar type represents non-fractional signed whole numeric values. Int can represent values between -(2^31) and 2^31 - 1.
 """
 
 
-String = TypeAliasType("String", str)
+String: TypeAlias = str
 """
 The `String` scalar type represents textual data, represented as UTF-8 character sequences. The String type is most often used by GraphQL to represent free-form human-readable text.
 """