feat: Add --enum-field-as-literal none option and respect user settings for TypedDict (#2691)

* feat: Add 'none' option for enum_field_as_literal and update related logic

* test: Skip tests for TypedDict enum_field_as_literal on incompatible black version

Author: koxudaxi

src/datamodel_code_generator/__init__.py

@@ -660,9 +660,9 @@ def get_header_and_first_line(csv_file: IO[str]) -> dict[str, Any]:
         reuse_model=reuse_model,
         reuse_scope=reuse_scope,
         shared_module_name=shared_module_name,
-        enum_field_as_literal=LiteralType.All
-        if output_model_type == DataModelType.TypingTypedDict
-        else enum_field_as_literal,
+        enum_field_as_literal=enum_field_as_literal
+        if enum_field_as_literal is not None
+        else (LiteralType.All if output_model_type == DataModelType.TypingTypedDict else None),
         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

src/datamodel_code_generator/arguments.py

@@ -343,7 +343,8 @@ def start_section(self, heading: str | None) -> None:
     "--enum-field-as-literal",
     help="Parse enum field as literal. "
     "all: all enum field type are Literal. "
-    "one: field type is Literal when an enum has only one possible value",
+    "one: field type is Literal when an enum has only one possible value. "
+    "none: always use Enum class (never convert to Literal)",
     choices=[lt.value for lt in LiteralType],
     default=None,
 )

src/datamodel_code_generator/parser/__init__.py

@@ -19,6 +19,7 @@ class LiteralType(Enum):
 
     All = "all"
     One = "one"
+    Off = "none"
 
 
 class DefaultPutDict(UserDict[TK, TV]):

src/datamodel_code_generator/parser/graphql.py

@@ -406,9 +406,11 @@ def parse_scalar(self, scalar_graphql_object: graphql.GraphQLScalarType) -> None
 
     def should_parse_enum_as_literal(self, obj: graphql.GraphQLEnumType) -> bool:
         """Determine if an enum should be parsed as a literal type."""
-        return self.enum_field_as_literal == LiteralType.All or (
-            self.enum_field_as_literal == LiteralType.One and len(obj.values) == 1
-        )
+        if self.enum_field_as_literal == LiteralType.All:
+            return True
+        if self.enum_field_as_literal == LiteralType.One:
+            return len(obj.values) == 1
+        return False
 
     def parse_enum(self, enum_object: graphql.GraphQLEnumType) -> None:
         """Parse a GraphQL enum type and add it to results."""

src/datamodel_code_generator/parser/jsonschema.py

@@ -767,9 +767,11 @@ def root_id(self, value: str | None) -> None:
 
     def should_parse_enum_as_literal(self, obj: JsonSchemaObject) -> bool:
         """Determine if an enum should be parsed as a literal type."""
-        return self.enum_field_as_literal == LiteralType.All or (
-            self.enum_field_as_literal == LiteralType.One and len(obj.enum) == 1
-        )
+        if self.enum_field_as_literal == LiteralType.All:
+            return True
+        if self.enum_field_as_literal == LiteralType.One:
+            return len(obj.enum) == 1
+        return False
 
     @classmethod
     def _extract_const_enum_from_combined(  # noqa: PLR0912

tests/data/expected/main/jsonschema/typed_dict_enum_literal_all.py

@@ -0,0 +1,12 @@
+# generated by datamodel-codegen:
+#   filename:  enum_literal_typed_dict.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Literal, NotRequired, TypedDict
+
+
+class EnumLiteralTypedDict(TypedDict):
+    status: NotRequired[Literal['active', 'inactive', 'pending']]
+    priority: NotRequired[Literal['high']]

tests/data/expected/main/jsonschema/typed_dict_enum_literal_none.py

@@ -0,0 +1,23 @@
+# generated by datamodel-codegen:
+#   filename:  enum_literal_typed_dict.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from enum import StrEnum
+from typing import NotRequired, TypedDict
+
+
+class Status(StrEnum):
+    active = 'active'
+    inactive = 'inactive'
+    pending = 'pending'
+
+
+class Priority(StrEnum):
+    high = 'high'
+
+
+class EnumLiteralTypedDict(TypedDict):
+    status: NotRequired[Status]
+    priority: NotRequired[Priority]

tests/data/expected/main/jsonschema/typed_dict_enum_literal_one.py

@@ -0,0 +1,19 @@
+# generated by datamodel-codegen:
+#   filename:  enum_literal_typed_dict.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from enum import StrEnum
+from typing import Literal, NotRequired, TypedDict
+
+
+class Status(StrEnum):
+    active = 'active'
+    inactive = 'inactive'
+    pending = 'pending'
+
+
+class EnumLiteralTypedDict(TypedDict):
+    status: NotRequired[Status]
+    priority: NotRequired[Literal['high']]

tests/data/jsonschema/enum_literal_typed_dict.json

@@ -0,0 +1,17 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "title": "EnumLiteralTypedDict",
+  "type": "object",
+  "properties": {
+    "status": {
+      "title": "Status",
+      "type": "string",
+      "enum": ["active", "inactive", "pending"]
+    },
+    "priority": {
+      "title": "Priority",
+      "type": "string",
+      "enum": ["high"]
+    }
+  }
+}

tests/main/jsonschema/test_main_jsonschema.py

@@ -2841,6 +2841,75 @@ def test_main_typed_dict_additional_properties(output_file: Path) -> None:
     )
 
 
+@pytest.mark.skipif(
+    black.__version__.split(".")[0] == "22",
+    reason="Installed black doesn't support Python version 3.11",
+)
+def test_main_typed_dict_enum_field_as_literal_none(output_file: Path) -> None:
+    """Test TypedDict with enum_field_as_literal=none."""
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / "enum_literal_typed_dict.json",
+        output_path=output_file,
+        input_file_type=None,
+        assert_func=assert_file_content,
+        expected_file="typed_dict_enum_literal_none.py",
+        extra_args=[
+            "--output-model-type",
+            "typing.TypedDict",
+            "--enum-field-as-literal",
+            "none",
+            "--target-python-version",
+            "3.11",
+        ],
+    )
+
+
+@pytest.mark.skipif(
+    black.__version__.split(".")[0] == "22",
+    reason="Installed black doesn't support Python version 3.11",
+)
+def test_main_typed_dict_enum_field_as_literal_one(output_file: Path) -> None:
+    """Test TypedDict with enum_field_as_literal=one."""
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / "enum_literal_typed_dict.json",
+        output_path=output_file,
+        input_file_type=None,
+        assert_func=assert_file_content,
+        expected_file="typed_dict_enum_literal_one.py",
+        extra_args=[
+            "--output-model-type",
+            "typing.TypedDict",
+            "--enum-field-as-literal",
+            "one",
+            "--target-python-version",
+            "3.11",
+        ],
+    )
+
+
+@pytest.mark.skipif(
+    black.__version__.split(".")[0] == "22",
+    reason="Installed black doesn't support Python version 3.11",
+)
+def test_main_typed_dict_enum_field_as_literal_all(output_file: Path) -> None:
+    """Test TypedDict with enum_field_as_literal=all."""
+    run_main_and_assert(
+        input_path=JSON_SCHEMA_DATA_PATH / "enum_literal_typed_dict.json",
+        output_path=output_file,
+        input_file_type=None,
+        assert_func=assert_file_content,
+        expected_file="typed_dict_enum_literal_all.py",
+        extra_args=[
+            "--output-model-type",
+            "typing.TypedDict",
+            "--enum-field-as-literal",
+            "all",
+            "--target-python-version",
+            "3.11",
+        ],
+    )
+
+
 def test_main_dataclass_const(output_file: Path) -> None:
     """Test main function writing to dataclass with const fields."""
     run_main_and_assert(