Support enum-field-as-literal in GraphQL parser (#2677)

* initial idea

* update dispatch

* clean unused

* fix required

* add tests

* tweak name / docstring

* cleanup unused

Author: siminn-arnorgj

src/datamodel_code_generator/parser/graphql.py

@@ -404,8 +404,40 @@ 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
+        )
+
     def parse_enum(self, enum_object: graphql.GraphQLEnumType) -> None:
         """Parse a GraphQL enum type and add it to results."""
+        if self.should_parse_enum_as_literal(enum_object):
+            return self.parse_enum_as_literal(enum_object)
+        return self.parse_enum_as_enum_class(enum_object)
+
+    def parse_enum_as_literal(self, enum_object: graphql.GraphQLEnumType) -> None:
+        """Parse enum values as a Literal type."""
+        data_type = self.data_type(literals=list(enum_object.values.keys()))
+        data_model_type = self._create_data_model(
+            model_type=self.data_model_root_type,
+            reference=self.references[enum_object.name],
+            fields=[
+                self.data_model_field_type(
+                    required=True,
+                    data_type=data_type,
+                )
+            ],
+            custom_base_class=self.base_class,
+            custom_template_dir=self.custom_template_dir,
+            extra_template_data=self.extra_template_data,
+            path=self.current_source_path,
+            description=enum_object.description,
+        )
+        self.results.append(data_model_type)
+
+    def parse_enum_as_enum_class(self, enum_object: graphql.GraphQLEnumType) -> None:
+        """Parse enum values as an Enum class."""
         enum_fields: list[DataModelFieldBase] = []
         exclude_field_names: set[str] = set()
 

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

@@ -0,0 +1,37 @@
+# generated by datamodel-codegen:
+#   filename:  enums.graphql
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import BaseModel
+from typing_extensions import TypeAlias
+
+Boolean: TypeAlias = bool
+"""
+The `Boolean` scalar type represents `true` or `false`.
+"""
+
+
+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.
+"""
+
+
+class Color(BaseModel):
+    __root__: Literal['BLUE', 'GREEN', 'RED']
+
+
+class EmployeeShiftStatus(BaseModel):
+    """
+    Employee shift status
+    """
+
+    __root__: Literal['NOT_ON_SHIFT', 'ON_SHIFT']
+
+
+class EnumWithOneField(BaseModel):
+    __root__: Literal['FIELD']

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

@@ -0,0 +1,41 @@
+# generated by datamodel-codegen:
+#   filename:  enums.graphql
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Literal
+
+from pydantic import BaseModel
+from typing_extensions import TypeAlias
+
+Boolean: TypeAlias = bool
+"""
+The `Boolean` scalar type represents `true` or `false`.
+"""
+
+
+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.
+"""
+
+
+class Color(Enum):
+    BLUE = 'BLUE'
+    GREEN = 'GREEN'
+    RED = 'RED'
+
+
+class EmployeeShiftStatus(Enum):
+    """
+    Employee shift status
+    """
+
+    NOT_ON_SHIFT = 'NOT_ON_SHIFT'
+    ON_SHIFT = 'ON_SHIFT'
+
+
+class EnumWithOneField(BaseModel):
+    __root__: Literal['FIELD']

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

@@ -0,0 +1,32 @@
+# generated by datamodel-codegen:
+#   filename:  enums.graphql
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Literal
+
+from typing_extensions import TypeAlias
+
+Boolean: TypeAlias = bool
+"""
+The `Boolean` scalar type represents `true` or `false`.
+"""
+
+
+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.
+"""
+
+
+Color: TypeAlias = Literal['BLUE', 'GREEN', 'RED']
+
+
+EmployeeShiftStatus: TypeAlias = Literal['NOT_ON_SHIFT', 'ON_SHIFT']
+"""
+Employee shift status
+"""
+
+
+EnumWithOneField: TypeAlias = Literal['FIELD']

tests/main/graphql/test_main_graphql.py

@@ -165,6 +165,42 @@ def test_main_graphql_specialized_enums(output_file: Path) -> None:
     )
 
 
+def test_main_graphql_enums_as_literals_all(output_file: Path) -> None:
+    """Test GraphQL code generation with all enum fields as literals."""
+    run_main_and_assert(
+        input_path=GRAPHQL_DATA_PATH / "enums.graphql",
+        output_path=output_file,
+        input_file_type="graphql",
+        assert_func=assert_file_content,
+        expected_file="enum_literals_all.py",
+        extra_args=["--enum-field-as-literal", "all"],
+    )
+
+
+def test_main_graphql_enums_as_literals_one(output_file: Path) -> None:
+    """Test GraphQL code generation with single-field enums as literals."""
+    run_main_and_assert(
+        input_path=GRAPHQL_DATA_PATH / "enums.graphql",
+        output_path=output_file,
+        input_file_type="graphql",
+        assert_func=assert_file_content,
+        expected_file="enum_literals_one.py",
+        extra_args=["--enum-field-as-literal", "one"],
+    )
+
+
+def test_main_graphql_enums_to_typed_dict(output_file: Path) -> None:
+    """Test GraphQL code generation paired with typing.TypedDict output which forces enums as literals."""
+    run_main_and_assert(
+        input_path=GRAPHQL_DATA_PATH / "enums.graphql",
+        output_path=output_file,
+        input_file_type="graphql",
+        assert_func=assert_file_content,
+        expected_file="enums_typed_dict.py",
+        extra_args=["--output-model-type", "typing.TypedDict"],
+    )
+
+
 @pytest.mark.skipif(
     black.__version__.split(".")[0] == "22",
     reason="Installed black doesn't support the old style",