Use Enum Members in Discriminators with Safe Imports (#2609)

* Feat: Add support for using enum member literals in discriminator fields

* Feat: Implement enum member imports for discriminator literals and add tests

* Refactor: Remove deprecated test for enum values in discriminator literals

* Refactor: Clean up code with pragma comments for better coverage tracking

* Refactor: Add pragma comments to improve coverage tracking in base.py

* Refactor: Simplify discriminator data type creation and enhance enum source resolution

Author: koxudaxi

README.md

@@ -381,6 +381,9 @@ Typing customization:
                         times.
   --use-annotated       Use typing.Annotated for Field(). Also, `--field-constraints` option
                         will be enabled.
+  --use-enum-values-in-discriminator
+                        Use enum member literals in discriminator fields instead of string
+                        literals
   --use-generic-container-types
                         Use generic container types for type hinting (typing.Sequence,
                         typing.Mapping). If `--use-standard-collections` option is set, then

docs/index.md

@@ -373,6 +373,9 @@ Typing customization:
                         times.
   --use-annotated       Use typing.Annotated for Field(). Also, `--field-constraints` option
                         will be enabled.
+  --use-enum-values-in-discriminator
+                        Use enum member literals in discriminator fields instead of string
+                        literals
   --use-generic-container-types
                         Use generic container types for type hinting (typing.Sequence,
                         typing.Mapping). If `--use-standard-collections` option is set, then

src/datamodel_code_generator/__init__.py

@@ -394,6 +394,7 @@ def generate(  # noqa: PLR0912, PLR0913, PLR0914, PLR0915
     encoding: str = "utf-8",
     enum_field_as_literal: LiteralType | None = None,
     use_one_literal_as_default: bool = False,
+    use_enum_values_in_discriminator: bool = False,
     set_default_enum_member: bool = False,
     use_subclass_enum: bool = False,
     use_specialized_enum: bool = True,
@@ -631,6 +632,7 @@ def get_header_and_first_line(csv_file: IO[str]) -> dict[str, Any]:
         if output_model_type == DataModelType.TypingTypedDict
         else enum_field_as_literal,
         use_one_literal_as_default=use_one_literal_as_default,
+        use_enum_values_in_discriminator=use_enum_values_in_discriminator,
         set_default_enum_member=True
         if output_model_type == DataModelType.DataclassesDataclass
         else set_default_enum_member,

src/datamodel_code_generator/__main__.py

@@ -398,6 +398,7 @@ def validate_all_exports_collision_strategy(cls, values: dict[str, Any]) -> dict
     encoding: str = DEFAULT_ENCODING
     enum_field_as_literal: Optional[LiteralType] = None  # noqa: UP045
     use_one_literal_as_default: bool = False
+    use_enum_values_in_discriminator: bool = False
     set_default_enum_member: bool = False
     use_subclass_enum: bool = False
     use_specialized_enum: bool = True
@@ -803,6 +804,7 @@ def main(args: Sequence[str] | None = None) -> Exit:  # noqa: PLR0911, PLR0912,
             encoding=config.encoding,
             enum_field_as_literal=config.enum_field_as_literal,
             use_one_literal_as_default=config.use_one_literal_as_default,
+            use_enum_values_in_discriminator=config.use_enum_values_in_discriminator,
             set_default_enum_member=config.set_default_enum_member,
             use_subclass_enum=config.use_subclass_enum,
             use_specialized_enum=config.use_specialized_enum,

src/datamodel_code_generator/arguments.py

@@ -373,6 +373,12 @@ def start_section(self, heading: str | None) -> None:
     action="store_true",
     default=None,
 )
+typing_options.add_argument(
+    "--use-enum-values-in-discriminator",
+    help="Use enum member literals in discriminator fields instead of string literals",
+    action="store_true",
+    default=None,
+)
 typing_options.add_argument(
     "--use-standard-collections",
     help="Use standard collections for type hinting (list, dict)",

src/datamodel_code_generator/parser/base.py

@@ -498,6 +498,7 @@ def __init__(  # noqa: PLR0913, PLR0915
         capitalise_enum_members: bool = False,
         keep_model_order: bool = False,
         use_one_literal_as_default: bool = False,
+        use_enum_values_in_discriminator: bool = False,
         known_third_party: list[str] | None = None,
         custom_formatters: list[str] | None = None,
         custom_formatters_kwargs: dict[str, Any] | None = None,
@@ -636,6 +637,7 @@ def __init__(  # noqa: PLR0913, PLR0915
         self.capitalise_enum_members = capitalise_enum_members
         self.keep_model_order = keep_model_order
         self.use_one_literal_as_default = use_one_literal_as_default
+        self.use_enum_values_in_discriminator = use_enum_values_in_discriminator
         self.known_third_party = known_third_party
         self.custom_formatter = custom_formatters
         self.custom_formatters_kwargs = custom_formatters_kwargs
@@ -912,6 +914,30 @@ def __extract_inherited_enum(cls, models: list[DataModel]) -> None:
                 )
                 models.remove(model)
 
+    def _create_discriminator_data_type(
+        self,
+        enum_source: Enum | None,
+        type_names: list[str],
+        discriminator_model: DataModel,
+        imports: Imports,
+    ) -> DataType:
+        """Create a data type for discriminator field, using enum literals if available."""
+        if enum_source:
+            enum_class_name = enum_source.reference.short_name
+            enum_member_literals: list[tuple[str, str]] = []
+            for value in type_names:
+                member = enum_source.find_member(value)
+                if member and member.field.name:
+                    enum_member_literals.append((enum_class_name, member.field.name))
+                else:  # pragma: no cover
+                    enum_member_literals.append((enum_class_name, value))
+            data_type = self.data_type(enum_member_literals=enum_member_literals)
+            if enum_source.module_path != discriminator_model.module_path:  # pragma: no cover
+                imports.append(Import.from_full_path(enum_source.name))
+        else:
+            data_type = self.data_type(literals=type_names)
+        return data_type
+
     def __apply_discriminator_type(  # noqa: PLR0912, PLR0915
         self,
         models: list[DataModel],
@@ -988,6 +1014,31 @@ def check_paths(
                         msg = f"Discriminator type is not found. {data_type.reference.path}"
                         raise RuntimeError(msg)
 
+                    enum_from_base: Enum | None = None
+                    if self.use_enum_values_in_discriminator:
+                        for base_class in discriminator_model.base_classes:
+                            if not base_class.reference or not base_class.reference.source:  # pragma: no cover
+                                continue
+                            base_model = base_class.reference.source
+                            if not isinstance(  # pragma: no cover
+                                base_model,
+                                (
+                                    pydantic_model.BaseModel,
+                                    pydantic_model_v2.BaseModel,
+                                    dataclass_model.DataClass,
+                                    msgspec_model.Struct,
+                                ),
+                            ):
+                                continue
+                            for base_field in base_model.fields:  # pragma: no branch
+                                if field_name not in {base_field.original_name, base_field.name}:  # pragma: no cover
+                                    continue
+                                enum_from_base = base_field.data_type.find_source(Enum)
+                                if enum_from_base:  # pragma: no branch
+                                    break
+                            if enum_from_base:  # pragma: no branch
+                                break
+
                     has_one_literal = False
                     for discriminator_field in discriminator_model.fields:
                         if field_name not in {discriminator_field.original_name, discriminator_field.name}:
@@ -1001,19 +1052,32 @@ def check_paths(
                                 discriminator_field.extras["is_classvar"] = True
                             # Found the discriminator field, no need to keep looking
                             break
+
+                        enum_source: Enum | None = None
+                        if self.use_enum_values_in_discriminator:
+                            enum_source = (  # pragma: no cover
+                                discriminator_field.data_type.find_source(Enum) or enum_from_base
+                            )
+
                         for field_data_type in discriminator_field.data_type.all_data_types:
                             if field_data_type.reference:  # pragma: no cover
                                 field_data_type.remove_reference()
-                        discriminator_field.data_type = self.data_type(literals=type_names)
+
+                        discriminator_field.data_type = self._create_discriminator_data_type(
+                            enum_source, type_names, discriminator_model, imports
+                        )
                         discriminator_field.data_type.parent = discriminator_field
                         discriminator_field.required = True
                         imports.append(discriminator_field.imports)
                         has_one_literal = True
                     if not has_one_literal:
+                        new_data_type = self._create_discriminator_data_type(
+                            enum_from_base, type_names, discriminator_model, imports
+                        )
                         discriminator_model.fields.append(
                             self.data_model_field_type(
                                 name=field_name,
-                                data_type=self.data_type(literals=type_names),
+                                data_type=new_data_type,
                                 required=True,
                                 alias=alias,
                             )

src/datamodel_code_generator/parser/graphql.py

@@ -167,6 +167,7 @@ def __init__(  # noqa: PLR0913
         capitalise_enum_members: bool = False,
         keep_model_order: bool = False,
         use_one_literal_as_default: bool = False,
+        use_enum_values_in_discriminator: bool = False,
         known_third_party: list[str] | None = None,
         custom_formatters: list[str] | None = None,
         custom_formatters_kwargs: dict[str, Any] | None = None,
@@ -222,6 +223,7 @@ def __init__(  # noqa: PLR0913
             encoding=encoding,
             enum_field_as_literal=enum_field_as_literal,
             use_one_literal_as_default=use_one_literal_as_default,
+            use_enum_values_in_discriminator=use_enum_values_in_discriminator,
             set_default_enum_member=set_default_enum_member,
             use_subclass_enum=use_subclass_enum,
             use_specialized_enum=use_specialized_enum,

src/datamodel_code_generator/parser/jsonschema.py

@@ -496,6 +496,7 @@ def __init__(  # noqa: PLR0913
         encoding: str = "utf-8",
         enum_field_as_literal: LiteralType | None = None,
         use_one_literal_as_default: bool = False,
+        use_enum_values_in_discriminator: bool = False,
         set_default_enum_member: bool = False,
         use_subclass_enum: bool = False,
         use_specialized_enum: bool = True,
@@ -585,6 +586,7 @@ def __init__(  # noqa: PLR0913
             encoding=encoding,
             enum_field_as_literal=enum_field_as_literal,
             use_one_literal_as_default=use_one_literal_as_default,
+            use_enum_values_in_discriminator=use_enum_values_in_discriminator,
             set_default_enum_member=set_default_enum_member,
             use_subclass_enum=use_subclass_enum,
             use_specialized_enum=use_specialized_enum,

src/datamodel_code_generator/parser/openapi.py

@@ -214,6 +214,7 @@ def __init__(  # noqa: PLR0913
         encoding: str = "utf-8",
         enum_field_as_literal: LiteralType | None = None,
         use_one_literal_as_default: bool = False,
+        use_enum_values_in_discriminator: bool = False,
         set_default_enum_member: bool = False,
         use_subclass_enum: bool = False,
         use_specialized_enum: bool = True,
@@ -305,6 +306,7 @@ def __init__(  # noqa: PLR0913
             encoding=encoding,
             enum_field_as_literal=enum_field_as_literal,
             use_one_literal_as_default=use_one_literal_as_default,
+            use_enum_values_in_discriminator=use_enum_values_in_discriminator,
             set_default_enum_member=set_default_enum_member,
             use_subclass_enum=use_subclass_enum,
             use_specialized_enum=use_specialized_enum,

src/datamodel_code_generator/types.py

@@ -54,6 +54,7 @@
 from datamodel_code_generator.util import PYDANTIC_V2, ConfigDict
 
 T = TypeVar("T")
+SourceT = TypeVar("SourceT")
 
 OPTIONAL = "Optional"
 OPTIONAL_PREFIX = f"{OPTIONAL}["
@@ -321,6 +322,7 @@ class Config:
     is_set: bool = False
     is_custom_type: bool = False
     literals: list[Union[StrictBool, StrictInt, StrictStr]] = []  # noqa: RUF012, UP007
+    enum_member_literals: list[tuple[str, str]] = []  # noqa: RUF012  # [(EnumClassName, member_name), ...]
     use_standard_collections: bool = False
     use_generic_container: bool = False
     use_union_operator: bool = False
@@ -406,6 +408,16 @@ def all_data_types(self) -> Iterator[DataType]:
             yield from data_type.all_data_types
         yield self
 
+    def find_source(self, source_type: type[SourceT]) -> SourceT | None:
+        """Find the first reference source matching the given type from all nested data types."""
+        for data_type in self.all_data_types:  # pragma: no branch
+            if not data_type.reference:  # pragma: no cover
+                continue
+            source = data_type.reference.source
+            if isinstance(source, source_type):  # pragma: no cover
+                return source
+        return None  # pragma: no cover
+
     @property
     def all_imports(self) -> Iterator[Import]:
         """Recursively yield all imports from nested DataTypes and self."""
@@ -424,7 +436,7 @@ def imports(self) -> Iterator[Import]:
         imports: tuple[tuple[bool, Import], ...] = (
             (self.is_optional and not self.use_union_operator, IMPORT_OPTIONAL),
             (len(self.data_types) > 1 and not self.use_union_operator, IMPORT_UNION),
-            (bool(self.literals), IMPORT_LITERAL),
+            (bool(self.literals) or bool(self.enum_member_literals), IMPORT_LITERAL),
         )
 
         if self.use_generic_container:
@@ -513,6 +525,9 @@ def type_hint(self) -> str:  # noqa: PLR0912, PLR0915
                     type_ = f"{UNION_PREFIX}{UNION_DELIMITER.join(data_types)}]"
             elif len(self.data_types) == 1:
                 type_ = self.data_types[0].type_hint
+            elif self.enum_member_literals:
+                parts = [f"{enum_class}.{member}" for enum_class, member in self.enum_member_literals]
+                type_ = f"{LITERAL}[{', '.join(parts)}]"
             elif self.literals:
                 type_ = f"{LITERAL}[{', '.join(repr(literal) for literal in self.literals)}]"
             elif self.reference: