support standard collections for type hinting (list, dict) (#263)

* support standard collections for type hinting (list, dict)

* fix types

README.md

@@ -65,19 +65,13 @@ $ pip install datamodel-code-generator[http]
 
 The `datamodel-codegen` command:
 ```
-usage: datamodel-codegen [-h] [--input INPUT]
-                         [--input-file-type {auto,openapi,jsonschema,json,yaml}]
-                         [--output OUTPUT] [--base-class BASE_CLASS]
-                         [--field-constraints] [--snake-case-field]
-                         [--strip-default-none]
-                         [--allow-population-by-field-name] [--use-default]
-                         [--force-optional] [--disable-timestamp]
-                         [--class-name CLASS_NAME]
-                         [--custom-template-dir CUSTOM_TEMPLATE_DIR]
-                         [--extra-template-data EXTRA_TEMPLATE_DATA]
-                         [--aliases ALIASES]
-                         [--target-python-version {3.6,3.7}] [--validation]
-                         [--debug] [--version]
+datamodel-codegen --help                                                                                master ✭ ✚ ✱ ◼
+usage: datamodel-codegen [-h] [--input INPUT] [--input-file-type {auto,openapi,jsonschema,json,yaml}] [--output OUTPUT]
+                         [--base-class BASE_CLASS] [--field-constraints] [--snake-case-field] [--strip-default-none]
+                         [--allow-population-by-field-name] [--use-default] [--force-optional] [--disable-timestamp]
+                         [--use-standard-collections] [--class-name CLASS_NAME]
+                         [--custom-template-dir CUSTOM_TEMPLATE_DIR] [--extra-template-data EXTRA_TEMPLATE_DATA]
+                         [--aliases ALIASES] [--target-python-version {3.6,3.7}] [--validation] [--debug] [--version]
 
 optional arguments:
   -h, --help            show this help message and exit
@@ -95,6 +89,8 @@ optional arguments:
   --use-default         Use default value even if a field is required
   --force-optional      Force optional for required fields
   --disable-timestamp   Disable timestamp on file headers
+  --use-standard-collections
+                        Use standard collections for type hinting (list, dict)
   --class-name CLASS_NAME
                         Set class name of root model
   --custom-template-dir CUSTOM_TEMPLATE_DIR

datamodel_code_generator/__init__.py

@@ -149,6 +149,7 @@ def generate(
     apply_default_values_for_required_fields: bool = False,
     force_optional_for_required_fields: bool = False,
     class_name: Optional[str] = None,
+    use_standard_collections: bool = False,
 ) -> None:
     input_text: Optional[str] = None
     if input_file_type == InputFileType.Auto:
@@ -205,6 +206,7 @@ def generate(
         apply_default_values_for_required_fields=apply_default_values_for_required_fields,
         force_optional_for_required_fields=force_optional_for_required_fields,
         class_name=class_name,
+        use_standard_collections=use_standard_collections,
     )
 
     with chdir(output):

datamodel_code_generator/__main__.py

@@ -107,6 +107,13 @@ def sig_int_handler(_: int, __: Any) -> None:  # pragma: no cover
     default=None,
 )
 
+arg_parser.add_argument(
+    '--use-standard-collections',
+    help='Use standard collections for type hinting (list, dict)',
+    action='store_true',
+    default=None,
+)
+
 arg_parser.add_argument(
     '--class-name', help='Set class name of root model', default=None,
 )
@@ -170,6 +177,7 @@ def validate_path(cls, value: Any) -> Optional[Path]:
     use_default: bool = False
     force_optional: bool = False
     class_name: Optional[str] = None
+    use_standard_collections: bool = False
 
     def merge_args(self, args: Namespace) -> None:
         for field_name in self.__fields__:
@@ -265,6 +273,7 @@ def main(args: Optional[Sequence[str]] = None) -> Exit:
             apply_default_values_for_required_fields=config.use_default,
             force_optional_for_required_fields=config.force_optional,
             class_name=config.class_name,
+            use_standard_collections=config.use_standard_collections,
         )
         return Exit.OK
     except InvalidClassNameError as e:

datamodel_code_generator/model/pydantic/types.py

@@ -76,6 +76,11 @@
     Types.any: DataType(type='Any', imports_=[IMPORT_ANY]),
 }
 
+standard_collections_type_map = {
+    **type_map,
+    Types.object: DataType(type='dict[str, Any]', imports_=[IMPORT_ANY,],),
+    Types.array: DataType(type='list[Any]', imports_=[IMPORT_ANY]),
+}
 kwargs_schema_to_model = {
     'exclusiveMinimum': 'gt',
     'minimum': 'ge',
@@ -109,9 +114,15 @@ def transform_kwargs(kwargs: Dict[str, Any], filter: Set[str]) -> Dict[str, str]
 
 
 class DataTypeManager(_DataTypeManager):
-    def __init__(self, python_version: PythonVersion = PythonVersion.PY_37):
-        super().__init__(python_version)
-        self.type_map: Dict[Types, DataType] = type_map
+    def __init__(
+        self,
+        python_version: PythonVersion = PythonVersion.PY_37,
+        use_standard_collections: bool = False,
+    ):
+        super().__init__(python_version, use_standard_collections)
+        self.type_map: Dict[
+            Types, DataType
+        ] = standard_collections_type_map if use_standard_collections else type_map
 
     def get_data_int_type(self, types: Types, **kwargs: Any) -> DataType:
         data_type_kwargs = transform_kwargs(kwargs, number_kwargs)

datamodel_code_generator/parser/base.py

@@ -166,9 +166,10 @@ def __init__(
         apply_default_values_for_required_fields: bool = False,
         force_optional_for_required_fields: bool = False,
         class_name: Optional[str] = None,
+        use_standard_collections: bool = False,
     ):
         self.data_type_manager: DataTypeManager = data_type_manager_type(
-            target_python_version
+            target_python_version, use_standard_collections
         )
         self.data_model_type: Type[DataModel] = data_model_type
         self.data_model_root_type: Type[DataModel] = data_model_root_type

datamodel_code_generator/parser/jsonschema.py

@@ -198,6 +198,7 @@ def __init__(
         apply_default_values_for_required_fields: bool = False,
         force_optional_for_required_fields: bool = False,
         class_name: Optional[str] = None,
+        use_standard_collections: bool = False,
     ):
         super().__init__(
             source=source,
@@ -219,6 +220,7 @@ def __init__(
             apply_default_values_for_required_fields=apply_default_values_for_required_fields,
             force_optional_for_required_fields=force_optional_for_required_fields,
             class_name=class_name,
+            use_standard_collections=use_standard_collections,
         )
 
         self.remote_object_cache: Dict[str, Dict[str, Any]] = {}

datamodel_code_generator/types.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from enum import Enum, auto
-from typing import Any, Dict, Iterator, List, Optional, Set, Type
+from typing import Any, Dict, Iterator, List, Optional, Set, Tuple, Type
 
 from pydantic import BaseModel
 
@@ -28,6 +28,7 @@ class DataType(BaseModel):
     is_optional: bool = False
     is_dict: bool = False
     is_list: bool = False
+    use_standard_collections: bool = False
 
     @classmethod
     def from_model_name(cls, model_name: str, is_list: bool = False) -> 'DataType':
@@ -51,12 +52,17 @@ def __init__(self, **values: Any) -> None:
         super().__init__(**values)
         if self.type and (self.reference or self.ref):
             self.unresolved_types.add(self.type)
-        for field, import_ in (
-            (self.is_list, IMPORT_LIST),
-            (self.is_dict, IMPORT_DICT),
+        imports: Tuple[Tuple[bool, Import], ...] = (
             (self.is_optional, IMPORT_OPTIONAL),
             (len(self.data_types) > 1, IMPORT_UNION),
-        ):
+        )
+        if not self.use_standard_collections:
+            imports = (
+                *imports,
+                (self.is_list, IMPORT_LIST),
+                (self.is_dict, IMPORT_DICT),
+            )
+        for field, import_ in imports:
             if field and import_ not in self.imports_:
                 self.imports_.append(import_)
         for data_type in self.data_types:
@@ -83,9 +89,17 @@ def type_hint(self) -> str:
                 # type_ = 'Any'
                 type_ = ''
         if self.is_list:
-            type_ = f'List[{type_}]' if type_ else 'List'
+            if self.use_standard_collections:
+                list_ = 'list'
+            else:
+                list_ = 'List'
+            type_ = f'{list_}[{type_}]' if type_ else list_
         if self.is_dict:
-            type_ = f'Dict[str, {type_}]' if type_ else 'Dict'
+            if self.use_standard_collections:
+                dict_ = 'dict'
+            else:
+                dict_ = 'Dict'
+            type_ = f'{dict_}[str, {type_}]' if type_ else 'dict_'
         if self.is_optional:
             type_ = f'Optional[{type_}]'
         if self.is_func:
@@ -103,6 +117,10 @@ class DataTypePy36(DataType):
     python_version: PythonVersion = PythonVersion.PY_36
 
 
+class DataTypeStandardCollections(DataType):
+    use_standard_collections: bool = True
+
+
 class Types(Enum):
     integer = auto()
     int32 = auto()
@@ -137,10 +155,16 @@ class Types(Enum):
 
 
 class DataTypeManager(ABC):
-    def __init__(self, python_version: PythonVersion = PythonVersion.PY_37) -> None:
+    def __init__(
+        self,
+        python_version: PythonVersion = PythonVersion.PY_37,
+        use_standard_collections: bool = False,
+    ) -> None:
         self.python_version = python_version
-        if python_version == PythonVersion.PY_36:
-            self.data_type: Type[DataType] = DataTypePy36
+        if use_standard_collections:
+            self.data_type: Type[DataType] = DataTypeStandardCollections
+        elif python_version == PythonVersion.PY_36:
+            self.data_type = DataTypePy36
         else:
             self.data_type = DataType
 

docs/index.md

@@ -32,19 +32,13 @@ $ pip install datamodel-code-generator[http]
 
 The `datamodel-codegen` command:
 ```
-usage: datamodel-codegen [-h] [--input INPUT]
-                         [--input-file-type {auto,openapi,jsonschema,json,yaml}]
-                         [--output OUTPUT] [--base-class BASE_CLASS]
-                         [--field-constraints] [--snake-case-field]
-                         [--strip-default-none]
-                         [--allow-population-by-field-name] [--use-default]
-                         [--force-optional] [--disable-timestamp]
-                         [--class-name CLASS_NAME]
-                         [--custom-template-dir CUSTOM_TEMPLATE_DIR]
-                         [--extra-template-data EXTRA_TEMPLATE_DATA]
-                         [--aliases ALIASES]
-                         [--target-python-version {3.6,3.7}] [--validation]
-                         [--debug] [--version]
+datamodel-codegen --help                                                                                master ✭ ✚ ✱ ◼
+usage: datamodel-codegen [-h] [--input INPUT] [--input-file-type {auto,openapi,jsonschema,json,yaml}] [--output OUTPUT]
+                         [--base-class BASE_CLASS] [--field-constraints] [--snake-case-field] [--strip-default-none]
+                         [--allow-population-by-field-name] [--use-default] [--force-optional] [--disable-timestamp]
+                         [--use-standard-collections] [--class-name CLASS_NAME]
+                         [--custom-template-dir CUSTOM_TEMPLATE_DIR] [--extra-template-data EXTRA_TEMPLATE_DATA]
+                         [--aliases ALIASES] [--target-python-version {3.6,3.7}] [--validation] [--debug] [--version]
 
 optional arguments:
   -h, --help            show this help message and exit
@@ -62,6 +56,8 @@ optional arguments:
   --use-default         Use default value even if a field is required
   --force-optional      Force optional for required fields
   --disable-timestamp   Disable timestamp on file headers
+  --use-standard-collections
+                        Use standard collections for type hinting (list, dict)
   --class-name CLASS_NAME
                         Set class name of root model
   --custom-template-dir CUSTOM_TEMPLATE_DIR

tests/data/expected/main/main_use_standard_collections/__init__.py

@@ -0,0 +1,34 @@
+# generated by datamodel-codegen:
+#   filename:  modular.yaml
+#   timestamp: 1985-10-26T08:21:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+from . import foo, models
+from .nested import foo as foo_1
+
+
+class Id(BaseModel):
+    __root__: str
+
+
+class Error(BaseModel):
+    code: int
+    message: str
+
+
+class Result(BaseModel):
+    event: Optional[models.Event] = None
+
+
+class Source(BaseModel):
+    country: Optional[str] = None
+
+
+class DifferentTea(BaseModel):
+    foo: Optional[foo.Tea] = None
+    nested: Optional[foo_1.Tea] = None

tests/data/expected/main/main_use_standard_collections/collections.py

@@ -0,0 +1,42 @@
+# generated by datamodel-codegen:
+#   filename:  modular.yaml
+#   timestamp: 1985-10-26T08:21:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import AnyUrl, BaseModel, Field
+
+from . import models
+
+
+class Pets(BaseModel):
+    __root__: list[models.Pet]
+
+
+class Users(BaseModel):
+    __root__: list[models.User]
+
+
+class Rules(BaseModel):
+    __root__: list[str]
+
+
+class Api(BaseModel):
+    apiKey: Optional[str] = Field(
+        None, description='To be used as a dataset parameter value'
+    )
+    apiVersionNumber: Optional[str] = Field(
+        None, description='To be used as a version parameter value'
+    )
+    apiUrl: Optional[AnyUrl] = Field(
+        None, description="The URL describing the dataset's fields"
+    )
+    apiDocumentationUrl: Optional[AnyUrl] = Field(
+        None, description='A URL to the API console for each API'
+    )
+
+
+class Apis(BaseModel):
+    __root__: list[Api]

tests/data/expected/main/main_use_standard_collections/foo/__init__.py

@@ -0,0 +1,20 @@
+# generated by datamodel-codegen:
+#   filename:  modular.yaml
+#   timestamp: 1985-10-26T08:21:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+from .. import Id
+
+
+class Tea(BaseModel):
+    flavour: Optional[str] = None
+    id: Optional[Id] = None
+
+
+class Cocoa(BaseModel):
+    quality: Optional[int] = None

tests/data/expected/main/main_use_standard_collections/foo/bar.py

@@ -0,0 +1,21 @@
+# generated by datamodel-codegen:
+#   filename:  modular.yaml
+#   timestamp: 1985-10-26T08:21:00+00:00
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+
+class Thing(BaseModel):
+    attributes: Optional[dict[str, Any]] = None
+
+
+class Thang(BaseModel):
+    attributes: Optional[list[dict[str, Any]]] = None
+
+
+class Clone(Thing):
+    pass