Add missing imports for constr and Dict[str, Any] (#69)

koxudaxi · Oct 10, 2019 · f5b43a6 · f5b43a6
1 parent 06abed6
commit f5b43a6
Show file tree

Hide file tree

Showing 8 changed files with 61 additions and 28 deletions.
diff --git a/datamodel_code_generator/imports.py b/datamodel_code_generator/imports.py
@@ -40,8 +40,10 @@ def append(self, imports: Union[Import, List[Import], None]) -> None:
                 self[import_.from_].add(import_.import_)
 
 
+IMPORT_ANY = Import(import_='Any', from_='typing')
 IMPORT_LIST = Import(import_='List', from_='typing')
 IMPORT_UNION = Import(import_='Union', from_='typing')
 IMPORT_OPTIONAL = Import(import_='Optional', from_='typing')
 IMPORT_ENUM = Import(import_='Enum', from_='enum')
 IMPORT_ANNOTATIONS = Import(from_='__future__', import_='annotations')
+IMPORT_CONSTR = Import(import_='constr', from_='pydantic')
diff --git a/datamodel_code_generator/model/base.py b/datamodel_code_generator/model/base.py
@@ -81,8 +81,8 @@ def __init__(self, **values: Any) -> None:
             self.alias = values['name']
         for data_type in self.data_types:
             self.unresolved_types.extend(data_type.unresolved_types)
-            if data_type.import_:
-                self.imports.append(data_type.import_)
+            if data_type.imports_:
+                self.imports.extend(data_type.imports_)
         self.type_hint = self._get_type_hint()
 
 

diff --git a/datamodel_code_generator/model/pydantic/types.py b/datamodel_code_generator/model/pydantic/types.py
@@ -1,6 +1,6 @@
 from typing import Any, Dict
 
-from datamodel_code_generator.imports import Import
+from datamodel_code_generator.imports import IMPORT_CONSTR, Import
 from datamodel_code_generator.types import DataType, Types
 
 type_map: Dict[Types, DataType] = {
@@ -14,40 +14,42 @@
     Types.string: DataType(type='str'),
     Types.byte: DataType(type='str'),  # base64 encoded string
     Types.binary: DataType(type='bytes'),
-    Types.date: DataType(type='date', import_=Import(from_='datetime', import_='date')),
+    Types.date: DataType(
+        type='date', imports_=[Import(from_='datetime', import_='date')]
+    ),
     Types.date_time: DataType(
-        type='datetime', import_=Import(from_='datetime', import_='datetime')
+        type='datetime', imports_=[Import(from_='datetime', import_='datetime')]
     ),
     Types.password: DataType(
-        type='SecretStr', import_=Import(from_='pydantic', import_='SecretStr')
+        type='SecretStr', imports_=[Import(from_='pydantic', import_='SecretStr')]
     ),
     Types.email: DataType(
-        type='EmailStr', import_=Import(from_='pydantic', import_='EmailStr')
+        type='EmailStr', imports_=[Import(from_='pydantic', import_='EmailStr')]
     ),
-    Types.uuid: DataType(type='UUID', import_=Import(from_='uuid', import_='UUID')),
+    Types.uuid: DataType(type='UUID', imports_=[Import(from_='uuid', import_='UUID')]),
     Types.uuid1: DataType(
-        type='UUID1', import_=Import(from_='pydantic', import_='UUID1')
+        type='UUID1', imports_=[Import(from_='pydantic', import_='UUID1')]
     ),
     Types.uuid2: DataType(
-        type='UUID2', import_=Import(from_='pydantic', import_='UUID2')
+        type='UUID2', imports_=[Import(from_='pydantic', import_='UUID2')]
     ),
     Types.uuid3: DataType(
-        type='UUID3', import_=Import(from_='pydantic', import_='UUID3')
+        type='UUID3', imports_=[Import(from_='pydantic', import_='UUID3')]
     ),
     Types.uuid4: DataType(
-        type='UUID4', import_=Import(from_='pydantic', import_='UUID4')
+        type='UUID4', imports_=[Import(from_='pydantic', import_='UUID4')]
     ),
     Types.uuid5: DataType(
-        type='UUID5', import_=Import(from_='pydantic', import_='UUID5')
+        type='UUID5', imports_=[Import(from_='pydantic', import_='UUID5')]
     ),
     Types.uri: DataType(
-        type='UrlStr', import_=Import(from_='pydantic', import_='UrlStr')
+        type='UrlStr', imports_=[Import(from_='pydantic', import_='UrlStr')]
     ),
     Types.ipv4: DataType(
-        type='IPv4Address', import_=Import(from_='pydantic', import_='IPv4Address')
+        type='IPv4Address', imports_=[Import(from_='pydantic', import_='IPv4Address')]
     ),
     Types.ipv6: DataType(
-        type='IPv6Address', import_=Import(from_='pydantic', import_='IPv6Address')
+        type='IPv6Address', imports_=[Import(from_='pydantic', import_='IPv6Address')]
     ),
     Types.boolean: DataType(type='bool'),
 }
@@ -106,7 +108,12 @@ def get_data_str_type(types: Types, **kwargs: Any) -> DataType:
     if kwargs.get('maxLength') is not None:
         data_type_kwargs['max_length'] = kwargs['maxLength']
     if data_type_kwargs:
-        return DataType(type='constr', is_func=True, kwargs=data_type_kwargs)
+        return DataType(
+            type='constr',
+            is_func=True,
+            kwargs=data_type_kwargs,
+            imports_=[IMPORT_CONSTR],
+        )
     return type_map[types]
 
 

diff --git a/datamodel_code_generator/parser/openapi.py b/datamodel_code_generator/parser/openapi.py
@@ -143,7 +143,15 @@ def parse_object_fields(self, obj: JsonSchemaObject) -> List[DataModelField]:
                         )
                     ]
                 else:
-                    field_types = [self.data_type(type='Dict[str, Any]')]
+                    field_types = [
+                        self.data_type(
+                            type='Dict[str, Any]',
+                            imports_=[
+                                Import(from_='typing', import_='Any'),
+                                Import(from_='typing', import_='Dict'),
+                            ],
+                        )
+                    ]
             elif field.enum:
                 enum = self.parse_enum(field_name, field)
                 field_types = [

diff --git a/datamodel_code_generator/types.py b/datamodel_code_generator/types.py
@@ -10,7 +10,7 @@ class DataType(BaseModel):
     type: str
     is_func: bool = False
     kwargs: Optional[Dict[str, Any]]
-    import_: Optional[Import]
+    imports_: Optional[List[Import]]
     python_version: PythonVersion = PythonVersion.PY_37
     unresolved_types: List[str] = []
     ref: bool = False

diff --git a/tests/model/pydantic/test_types.py b/tests/model/pydantic/test_types.py
@@ -1,4 +1,5 @@
 import pytest
+from datamodel_code_generator.imports import IMPORT_CONSTR
 from datamodel_code_generator.model.pydantic.types import (
     get_data_float_type,
     get_data_int_type,
@@ -89,17 +90,32 @@ def test_get_data_float_type(types, params, data_type):
         (
             Types.string,
             {'pattern': '^abc'},
-            DataType(type='constr', is_func=True, kwargs={'regex': '^abc'}),
+            DataType(
+                type='constr',
+                is_func=True,
+                kwargs={'regex': '^abc'},
+                imports_=[IMPORT_CONSTR],
+            ),
         ),
         (
             Types.string,
             {'minLength': 10},
-            DataType(type='constr', is_func=True, kwargs={'min_length': 10}),
+            DataType(
+                type='constr',
+                is_func=True,
+                kwargs={'min_length': 10},
+                imports_=[IMPORT_CONSTR],
+            ),
         ),
         (
             Types.string,
             {'maxLength': 10},
-            DataType(type='constr', is_func=True, kwargs={'max_length': 10}),
+            DataType(
+                type='constr',
+                is_func=True,
+                kwargs={'max_length': 10},
+                imports_=[IMPORT_CONSTR],
+            ),
         ),
     ],
 )

diff --git a/tests/parser/test_openapi.py b/tests/parser/test_openapi.py
@@ -1,6 +1,6 @@
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from typing import Optional
+from typing import List, Optional
 
 import pytest
 from datamodel_code_generator import PythonVersion
@@ -51,14 +51,14 @@ def render(self) -> str:
 )
 def test_get_data_type(schema_type, schema_format, result_type, from_, import_):
     if from_ and import_:
-        import_obj: Optional[Import] = Import(from_=from_, import_=import_)
+        imports_: Optional[List[Import]] = [Import(from_=from_, import_=import_)]
     else:
-        import_obj = None
+        imports_ = None
 
     parser = OpenAPIParser(BaseModel, CustomRootType)
     assert parser.get_data_type(
         JsonSchemaObject(type=schema_type, format=schema_format)
-    ) == DataType(type=result_type, import_=import_obj)
+    ) == DataType(type=result_type, imports_=imports_)
 
 
 def test_get_data_type_invalid_obj():
@@ -1001,7 +1001,7 @@ class apis(BaseModel):
                 ): '''\
 from __future__ import annotations
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from pydantic import BaseModel
 

diff --git a/tests/test_main.py b/tests/test_main.py
@@ -385,7 +385,7 @@ class apis(BaseModel):
 
 from __future__ import annotations
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from pydantic import BaseModel