add --capitalise-enum-members option (#970)

* add --capitalise-enum-members option

* Add test files

README.md

@@ -100,6 +100,7 @@ usage: datamodel-codegen [-h] [--input INPUT] [--url URL]
                          [--collapse-root-models] [--enum-field-as-literal {all,one}]
                          [--set-default-enum-member]
                          [--empty-enum-field-name EMPTY_ENUM_FIELD_NAME]
+                         [--capitalise-enum-members]
                          [--special-field-name-prefix SPECIAL_FIELD_NAME_PREFIX]
                          [--use-subclass-enum]
                          [--class-name CLASS_NAME] [--use-title-as-name]
@@ -180,6 +181,8 @@ optional arguments:
                         Set enum members as default values for enum field
   --empty-enum-field-name EMPTY_ENUM_FIELD_NAME
                         Set field name when enum value is empty (default: `_`)
+  --capitalise-enum-members
+                        Capitalize field names on enum
   --special-field-name-prefix SPECIAL_FIELD_NAME_PREFIX
   --use-subclass-enum   Define Enum class as subclass with field type when enum has
                         type (int, float, bytes, str)

datamodel_code_generator/__init__.py

@@ -248,6 +248,7 @@ def generate(
     use_union_operator: bool = False,
     collapse_root_models: bool = False,
     special_field_name_prefix: Optional[str] = None,
+    capitalise_enum_members: bool = False,
 ) -> None:
     remote_text_cache: DefaultPutDict[str, str] = DefaultPutDict()
     if isinstance(input_, str):
@@ -375,6 +376,7 @@ def get_header_and_first_line(csv_file: IO[str]) -> Dict[str, Any]:
         use_union_operator=use_union_operator,
         collapse_root_models=collapse_root_models,
         special_field_name_prefix=special_field_name_prefix,
+        capitalise_enum_members=capitalise_enum_members,
         **kwargs,
     )
 

datamodel_code_generator/__main__.py

@@ -297,6 +297,14 @@ def sig_int_handler(_: int, __: Any) -> None:  # pragma: no cover
     default=None,
 )
 
+
+arg_parser.add_argument(
+    '--capitalise-enum-members',
+    help='Capitalize field names on enum',
+    action='store_true',
+    default=None,
+)
+
 arg_parser.add_argument(
     '--special-field-name-prefix',
     help='Set field name prefix when first character can\'t be used as Python field name (default:  `field`)',
@@ -513,6 +521,7 @@ def _validate_use_union_operator(cls, values: Dict[str, Any]) -> Dict[str, Any]:
     use_double_quotes: bool = False
     collapse_root_models: bool = False
     special_field_name_prefix: Optional[str] = None
+    capitalise_enum_members: bool = False
 
     def merge_args(self, args: Namespace) -> None:
         set_args = {
@@ -661,6 +670,7 @@ def main(args: Optional[Sequence[str]] = None) -> Exit:
             collapse_root_models=config.collapse_root_models,
             use_union_operator=config.use_union_operator,
             special_field_name_prefix=config.special_field_name_prefix,
+            capitalise_enum_members=config.capitalise_enum_members,
         )
         return Exit.OK
     except InvalidClassNameError as e:

datamodel_code_generator/parser/base.py

@@ -313,6 +313,7 @@ def __init__(
         allow_responses_without_content: bool = False,
         collapse_root_models: bool = False,
         special_field_name_prefix: Optional[str] = None,
+        capitalise_enum_members: bool = False,
     ):
         self.data_type_manager: DataTypeManager = data_type_manager_type(
             python_version=target_python_version,
@@ -398,6 +399,7 @@ def __init__(
             base_path=self.base_path,
             original_field_name_delimiter=original_field_name_delimiter,
             special_field_name_prefix=special_field_name_prefix,
+            capitalise_enum_members=capitalise_enum_members,
         )
         self.class_name: Optional[str] = class_name
         self.wrap_string_literal: Optional[bool] = wrap_string_literal
@@ -414,6 +416,7 @@ def __init__(
         self.use_double_quotes = use_double_quotes
         self.allow_responses_without_content = allow_responses_without_content
         self.collapse_root_models = collapse_root_models
+        self.capitalise_enum_members = capitalise_enum_members
 
     @property
     def iter_source(self) -> Iterator[Source]:

datamodel_code_generator/parser/jsonschema.py

@@ -347,6 +347,7 @@ def __init__(
         allow_responses_without_content: bool = False,
         collapse_root_models: bool = False,
         special_field_name_prefix: Optional[str] = None,
+        capitalise_enum_members: bool = False,
     ):
         super().__init__(
             source=source,
@@ -400,6 +401,7 @@ def __init__(
             allow_responses_without_content=allow_responses_without_content,
             collapse_root_models=collapse_root_models,
             special_field_name_prefix=special_field_name_prefix,
+            capitalise_enum_members=capitalise_enum_members,
         )
 
         self.remote_object_cache: DefaultPutDict[str, Dict[str, Any]] = DefaultPutDict()

datamodel_code_generator/parser/openapi.py

@@ -196,6 +196,7 @@ def __init__(
         allow_responses_without_content: bool = False,
         collapse_root_models: bool = False,
         special_field_name_prefix: Optional[str] = None,
+        capitalise_enum_members: bool = False,
     ):
         super().__init__(
             source=source,
@@ -249,6 +250,7 @@ def __init__(
             allow_responses_without_content=allow_responses_without_content,
             collapse_root_models=collapse_root_models,
             special_field_name_prefix=special_field_name_prefix,
+            capitalise_enum_members=capitalise_enum_members,
         )
         self.open_api_scopes: List[OpenAPIScope] = openapi_scopes or [
             OpenAPIScope.Schemas

datamodel_code_generator/reference.py

@@ -143,6 +143,7 @@ def __init__(
         empty_field_name: Optional[str] = None,
         original_delimiter: Optional[str] = None,
         special_field_name_prefix: Optional[str] = None,
+        capitalise_enum_members: bool = False,
     ):
         self.aliases: Mapping[str, str] = {} if aliases is None else {**aliases}
         self.empty_field_name: str = empty_field_name or '_'
@@ -151,6 +152,7 @@ def __init__(
         self.special_field_name_prefix: Optional[str] = (
             'field' if special_field_name_prefix is None else special_field_name_prefix
         )
+        self.capitalise_enum_members: bool = capitalise_enum_members
 
     @classmethod
     def _validate_field_name(cls, field_name: str) -> bool:
@@ -183,12 +185,21 @@ def get_valid_name(
         # causes pydantic to consider it as private
         if name.startswith('_'):
             name = f'{self.special_field_name_prefix}{name}'
-        if self.snake_case_field and not ignore_snake_case_field:
+        if (
+            self.capitalise_enum_members
+            or self.snake_case_field
+            and not ignore_snake_case_field
+        ):
             name = camel_to_snake(name)
         count = 1
         if iskeyword(name) or not self._validate_field_name(name):
             name += '_'
-        new_name = snake_to_upper_camel(name) if upper_camel else name
+        if upper_camel:
+            new_name = snake_to_upper_camel(name)
+        elif self.capitalise_enum_members:
+            new_name = name.upper()
+        else:
+            new_name = name
         while (
             not (new_name.isidentifier() or not self._validate_field_name(new_name))
             or iskeyword(new_name)
@@ -276,6 +287,7 @@ def __init__(
         ] = None,
         original_field_name_delimiter: Optional[str] = None,
         special_field_name_prefix: Optional[str] = None,
+        capitalise_enum_members: bool = False,
     ) -> None:
         self.references: Dict[str, Reference] = {}
         self._current_root: Sequence[str] = []
@@ -301,6 +313,7 @@ def __init__(
                 empty_field_name=empty_field_name,
                 original_delimiter=original_field_name_delimiter,
                 special_field_name_prefix=special_field_name_prefix,
+                capitalise_enum_members=capitalise_enum_members,
             )
             for k, v in merged_field_name_resolver_classes.items()
         }

docs/index.md

@@ -59,6 +59,7 @@ usage: datamodel-codegen [-h] [--input INPUT] [--url URL]
                          [--collapse-root-models] [--enum-field-as-literal {all,one}]
                          [--set-default-enum-member]
                          [--empty-enum-field-name EMPTY_ENUM_FIELD_NAME]
+                         [--capitalise-enum-members]
                          [--special-field-name-prefix SPECIAL_FIELD_NAME_PREFIX]
                          [--use-subclass-enum]
                          [--class-name CLASS_NAME] [--use-title-as-name]
@@ -141,6 +142,8 @@ optional arguments:
   --set-default-enum-member
                         Set enum members as default values for enum field
   --empty-enum-field-name EMPTY_ENUM_FIELD_NAME
+  --capitalise-enum-members
+                        Capitalize field names on enum
   --special-field-name-prefix SPECIAL_FIELD_NAME_PREFIX
                         Set field name when enum value is empty (default: `_`)
   --use-subclass-enum   Define Enum class as subclass with field type when enum has

tests/data/expected/main/main_json_capitalise_enum_members/output.py

@@ -0,0 +1,14 @@
+# generated by datamodel-codegen:
+#   filename:  many_case_enum.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from enum import Enum
+
+
+class Model(Enum):
+    SNAKE_CASE = 'snake_case'
+    CAP_CASE = 'CAP_CASE'
+    CAMEL_CASE = 'CamelCase'
+    UPPERCASE = 'UPPERCASE'

tests/data/jsonschema/many_case_enum.json

@@ -0,0 +1,10 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "type": "string",
+  "enum": [
+    "snake_case",
+    "CAP_CASE",
+    "CamelCase",
+    "UPPERCASE"
+  ]
+}

tests/test_main.py

@@ -1755,6 +1755,32 @@ def test_main_json_reuse_enum():
         main()
 
 
+@freeze_time('2019-07-26')
+def test_main_json_capitalise_enum_members():
+    with TemporaryDirectory() as output_dir:
+        output_file: Path = Path(output_dir) / 'output.py'
+        return_code: Exit = main(
+            [
+                '--input',
+                str(JSON_SCHEMA_DATA_PATH / 'many_case_enum.json'),
+                '--output',
+                str(output_file),
+                '--input-file-type',
+                'jsonschema',
+                '--capitalise-enum-members',
+            ]
+        )
+        assert return_code == Exit.OK
+        assert (
+            output_file.read_text()
+            == (
+                EXPECTED_MAIN_PATH / 'main_json_capitalise_enum_members' / 'output.py'
+            ).read_text()
+        )
+    with pytest.raises(SystemExit):
+        main()
+
+
 @freeze_time('2019-07-26')
 def test_main_openapi_datetime():
     with TemporaryDirectory() as output_dir: