Support --original-field-name-delimiter option (#776)

* Support --original-field-name-delimiter option * add validator for the option
koxudaxi · May 27, 2022 · ce9bed7 · ce9bed7
1 parent 7c7e35d
commit ce9bed7
Show file tree

Hide file tree

Showing 11 changed files with 115 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -79,6 +79,7 @@ usage: datamodel-codegen [-h] [--input INPUT] [--url URL]
                          [--use_non_positive_negative_number_constrained_types]
                          [--field-extra-keys FIELD_EXTRA_KEYS [FIELD_EXTRA_KEYS ...]]
                          [--field-include-all-keys] [--snake-case-field]
+                         [--original-field-name-delimiter ORIGINAL_FIELD_NAME_DELIMITER]
                          [--strip-default-none]
                          [--disable-appending-item-suffix]
                          [--allow-population-by-field-name]
@@ -128,6 +129,9 @@ optional arguments:
   --field-include-all-keys
                         Add all keys to field parameters
   --snake-case-field    Change camel-case field name to snake-case
+  --original-field-name-delimiter ORIGINAL_FIELD_NAME_DELIMITER
+                        Set delimiter to convert to snake case. This option only
+                        can be used with --snake-case-field (default: `_` )
   --strip-default-none  Strip default None on fields
   --disable-appending-item-suffix
                         Disable appending `Item` suffix to model name in an

diff --git a/datamodel_code_generator/__init__.py b/datamodel_code_generator/__init__.py
@@ -238,6 +238,7 @@ def generate(
     http_ignore_tls: bool = False,
     use_annotated: bool = False,
     use_non_positive_negative_number_constrained_types: bool = False,
+    original_field_name_delimiter: Optional[str] = None,
 ) -> None:
     remote_text_cache: DefaultPutDict[str, str] = DefaultPutDict()
     if isinstance(input_, str):
@@ -358,6 +359,7 @@ def get_header_and_first_line(csv_file: IO[str]) -> Dict[str, Any]:
         http_ignore_tls=http_ignore_tls,
         use_annotated=use_annotated,
         use_non_positive_negative_number_constrained_types=use_non_positive_negative_number_constrained_types,
+        original_field_name_delimiter=original_field_name_delimiter,
         **kwargs,
     )
 

diff --git a/datamodel_code_generator/__main__.py b/datamodel_code_generator/__main__.py
@@ -146,6 +146,13 @@ def sig_int_handler(_: int, __: Any) -> None:  # pragma: no cover
     action='store_true',
     default=None,
 )
+
+arg_parser.add_argument(
+    '--original-field-name-delimiter',
+    help='Set delimiter to convert to snake case. This option only can be used with --snake-case-field (default: `_` )',
+    default=None,
+)
+
 arg_parser.add_argument(
     '--strip-default-none',
     help='Strip default None on fields',
@@ -358,6 +365,17 @@ def validate_use_generic_container_types(
                 )
         return values
 
+    @root_validator
+    def validate_original_field_name_delimiter(
+        cls, values: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        if values.get('original_field_name_delimiter') is not None:
+            if not values.get('snake_case_field'):
+                raise Error(
+                    "`--original-field-name-delimiter` can not be used without `--snake-case-field`."
+                )
+        return values
+
     # Pydantic 1.5.1 doesn't support each_item=True correctly
     @validator('http_headers', pre=True)
     def validate_http_headers(cls, value: Any) -> Optional[List[Tuple[str, str]]]:
@@ -427,6 +445,7 @@ def _validate_use_annotated(cls, values: Dict[str, Any]) -> Dict[str, Any]:
     http_ignore_tls: bool = False
     use_annotated: bool = False
     use_non_positive_negative_number_constrained_types: bool = False
+    original_field_name_delimiter: Optional[str] = None
 
     def merge_args(self, args: Namespace) -> None:
         set_args = {
@@ -559,6 +578,7 @@ def main(args: Optional[Sequence[str]] = None) -> Exit:
             http_ignore_tls=config.http_ignore_tls,
             use_annotated=config.use_annotated,
             use_non_positive_negative_number_constrained_types=config.use_non_positive_negative_number_constrained_types,
+            original_field_name_delimiter=config.original_field_name_delimiter,
         )
         return Exit.OK
     except InvalidClassNameError as e:

diff --git a/datamodel_code_generator/parser/base.py b/datamodel_code_generator/parser/base.py
@@ -299,6 +299,7 @@ def __init__(
         http_ignore_tls: bool = False,
         use_annotated: bool = False,
         use_non_positive_negative_number_constrained_types: bool = False,
+        original_field_name_delimiter: Optional[str] = None,
     ):
         self.data_type_manager: DataTypeManager = data_type_manager_type(
             target_python_version,
@@ -376,6 +377,7 @@ def __init__(
             snake_case_field=snake_case_field,
             custom_class_name_generator=custom_class_name_generator,
             base_path=self.base_path,
+            original_field_name_delimiter=original_field_name_delimiter,
         )
         self.class_name: Optional[str] = class_name
         self.wrap_string_literal: Optional[bool] = wrap_string_literal

diff --git a/datamodel_code_generator/parser/jsonschema.py b/datamodel_code_generator/parser/jsonschema.py
@@ -328,6 +328,7 @@ def __init__(
         http_ignore_tls: bool = False,
         use_annotated: bool = False,
         use_non_positive_negative_number_constrained_types: bool = False,
+        original_field_name_delimiter: Optional[str] = None,
     ):
         super().__init__(
             source=source,
@@ -373,6 +374,7 @@ def __init__(
             http_ignore_tls=http_ignore_tls,
             use_annotated=use_annotated,
             use_non_positive_negative_number_constrained_types=use_non_positive_negative_number_constrained_types,
+            original_field_name_delimiter=original_field_name_delimiter,
         )
 
         self.remote_object_cache: DefaultPutDict[str, Dict[str, Any]] = DefaultPutDict()

diff --git a/datamodel_code_generator/parser/openapi.py b/datamodel_code_generator/parser/openapi.py
@@ -183,6 +183,7 @@ def __init__(
         http_ignore_tls: bool = False,
         use_annotated: bool = False,
         use_non_positive_negative_number_constrained_types: bool = False,
+        original_field_name_delimiter: Optional[str] = None,
     ):
         super().__init__(
             source=source,
@@ -228,6 +229,7 @@ def __init__(
             http_ignore_tls=http_ignore_tls,
             use_annotated=use_annotated,
             use_non_positive_negative_number_constrained_types=use_non_positive_negative_number_constrained_types,
+            original_field_name_delimiter=original_field_name_delimiter,
         )
         self.open_api_scopes: List[OpenAPIScope] = openapi_scopes or [
             OpenAPIScope.Schemas

diff --git a/datamodel_code_generator/reference.py b/datamodel_code_generator/reference.py
@@ -131,10 +131,12 @@ def __init__(
         aliases: Optional[Mapping[str, str]] = None,
         snake_case_field: bool = False,
         empty_field_name: Optional[str] = None,
+        original_delimiter: Optional[str] = None,
     ):
         self.aliases: Mapping[str, str] = {} if aliases is None else {**aliases}
         self.empty_field_name: str = empty_field_name or '_'
         self.snake_case_field = snake_case_field
+        self.original_delimiter: Optional[str] = original_delimiter
 
     @classmethod
     def _validate_field_name(cls, field_name: str) -> bool:
@@ -149,6 +151,10 @@ def get_valid_name(
             name = self.empty_field_name
         if name[0] == '#':
             name = name[1:] or self.empty_field_name
+
+        if self.snake_case_field and self.original_delimiter is not None:
+            name = snake_to_upper_camel(name, delimiter=self.original_delimiter)
+
         # TODO: when first character is a number
         name = re.sub(r'\W', '_', name)
         if name[0].isnumeric():
@@ -230,6 +236,7 @@ def __init__(
         field_name_resolver_classes: Optional[
             Dict[ModelType, Type[FieldNameResolver]]
         ] = None,
+        original_field_name_delimiter: Optional[str] = None,
     ) -> None:
         self.references: Dict[str, Reference] = {}
         self._current_root: Sequence[str] = []
@@ -253,6 +260,7 @@ def __init__(
                 aliases=aliases,
                 snake_case_field=snake_case_field,
                 empty_field_name=empty_field_name,
+                original_delimiter=original_field_name_delimiter,
             )
             for k, v in merged_field_name_resolver_classes.items()
         }
@@ -583,13 +591,13 @@ def get_singular_name(name: str, suffix: str = SINGULAR_NAME_SUFFIX) -> str:
 
 
 @lru_cache()
-def snake_to_upper_camel(word: str) -> str:
+def snake_to_upper_camel(word: str, delimiter: str = '_') -> str:
     prefix = ''
-    if word.startswith('_'):
+    if word.startswith(delimiter):
         prefix = '_'
         word = word[1:]
 
-    return prefix + ''.join(x[0].upper() + x[1:] for x in word.split('_') if x)
+    return prefix + ''.join(x[0].upper() + x[1:] for x in word.split(delimiter) if x)
 
 
 def is_url(ref: str) -> bool:

diff --git a/docs/index.md b/docs/index.md
@@ -45,6 +45,7 @@ usage: datamodel-codegen [-h] [--input INPUT] [--url URL]
                          [--use_non_positive_negative_number_constrained_types]
                          [--field-extra-keys FIELD_EXTRA_KEYS [FIELD_EXTRA_KEYS ...]]
                          [--field-include-all-keys] [--snake-case-field]
+                         [--original-field-name-delimiter ORIGINAL_FIELD_NAME_DELIMITER]
                          [--strip-default-none]
                          [--disable-appending-item-suffix]
                          [--allow-population-by-field-name]
@@ -94,6 +95,9 @@ optional arguments:
   --field-include-all-keys
                         Add all keys to field parameters
   --snake-case-field    Change camel-case field name to snake-case
+  --original-field-name-delimiter ORIGINAL_FIELD_NAME_DELIMITER
+                        Set delimiter to convert to snake case. This option only
+                        can be used with --snake-case-field (default: `_` )
   --strip-default-none  Strip default None on fields
   --disable-appending-item-suffix
                         Disable appending `Item` suffix to model name in an

diff --git a/tests/data/expected/main/main_space_field_enum_snake_case_field/output.py b/tests/data/expected/main/main_space_field_enum_snake_case_field/output.py
@@ -0,0 +1,11 @@
+# generated by datamodel-codegen:
+#   filename:  space_field_enum.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from enum import Enum
+
+
+class Model(Enum):
+    space_field = 'Space Field'
diff --git a/tests/data/jsonschema/space_field_enum.json b/tests/data/jsonschema/space_field_enum.json
@@ -0,0 +1,7 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "type": "string",
+  "enum": [
+    "Space Field"
+  ]
+}
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -1468,6 +1468,25 @@ def test_main_use_generic_container_types_py36(capsys) -> None:
     )
 
 
+def test_main_original_field_name_delimiter_without_snake_case_field(capsys) -> None:
+    input_filename = OPEN_API_DATA_PATH / 'modular.yaml'
+
+    return_code: Exit = main(
+        [
+            '--input',
+            str(input_filename),
+            '--original-field-name-delimiter',
+            '-',
+        ]
+    )
+    captured = capsys.readouterr()
+    assert return_code == Exit.ERROR
+    assert (
+        captured.err
+        == '`--original-field-name-delimiter` can not be used without `--snake-case-field`.\n'
+    )
+
+
 @freeze_time('2019-07-26')
 def test_main_external_definitions():
     with TemporaryDirectory() as output_dir:
@@ -2152,6 +2171,37 @@ def test_simple_json_snake_case_field():
         main()
 
 
+@freeze_time('2019-07-26')
+def test_main_space_field_enum_snake_case_field():
+    with TemporaryDirectory() as output_dir:
+        output_file: Path = Path(output_dir) / 'output.py'
+        with chdir(JSON_SCHEMA_DATA_PATH / 'space_field_enum.json'):
+            return_code: Exit = main(
+                [
+                    '--input',
+                    'space_field_enum.json',
+                    '--output',
+                    str(output_file),
+                    '--input-file-type',
+                    'jsonschema',
+                    '--snake-case-field',
+                    '--original-field-name-delimiter',
+                    ' ',
+                ]
+            )
+            assert return_code == Exit.OK
+            assert (
+                output_file.read_text()
+                == (
+                    EXPECTED_MAIN_PATH
+                    / 'main_space_field_enum_snake_case_field'
+                    / 'output.py'
+                ).read_text()
+            )
+    with pytest.raises(SystemExit):
+        main()
+
+
 @freeze_time('2019-07-26')
 def test_main_all_of_ref():
     with TemporaryDirectory() as output_dir: