[FIX] types requiring imports in alloff, conint, confloat (#84)

* types requiring imports in alloff, conint, confloat * format, lint
koxudaxi · Dec 13, 2019 · 2c9e777 · 2c9e777
1 parent ad35e4c
commit 2c9e777
Show file tree

Hide file tree

Showing 9 changed files with 822 additions and 15 deletions.
diff --git a/datamodel_code_generator/imports.py b/datamodel_code_generator/imports.py
@@ -47,3 +47,5 @@ def append(self, imports: Union[Import, List[Import], None]) -> None:
 IMPORT_ENUM = Import(import_='Enum', from_='enum')
 IMPORT_ANNOTATIONS = Import(from_='__future__', import_='annotations')
 IMPORT_CONSTR = Import(import_='constr', from_='pydantic')
+IMPORT_CONINT = Import(import_='conint', from_='pydantic')
+IMPORT_CONFLOAT = Import(import_='confloat', from_='pydantic')
diff --git a/datamodel_code_generator/model/pydantic/base_model.py b/datamodel_code_generator/model/pydantic/base_model.py
@@ -27,6 +27,7 @@ def __init__(
         extra_template_data: Optional[DefaultDict[str, Any]] = None,
         auto_import: bool = True,
         reference_classes: Optional[List[str]] = None,
+        imports: Optional[List[Import]] = None,
     ):
 
         super().__init__(
@@ -39,6 +40,7 @@ def __init__(
             extra_template_data=extra_template_data,
             auto_import=auto_import,
             reference_classes=reference_classes,
+            imports=imports,
         )
 
         if 'additionalProperties' in self.extra_template_data:

diff --git a/datamodel_code_generator/model/pydantic/types.py b/datamodel_code_generator/model/pydantic/types.py
@@ -1,6 +1,11 @@
 from typing import Any, Dict
 
-from datamodel_code_generator.imports import IMPORT_CONSTR, Import
+from datamodel_code_generator.imports import (
+    IMPORT_CONFLOAT,
+    IMPORT_CONINT,
+    IMPORT_CONSTR,
+    Import,
+)
 from datamodel_code_generator.types import DataType, Types
 
 type_map: Dict[Types, DataType] = {
@@ -80,7 +85,12 @@ def get_data_int_type(types: Types, **kwargs: Any) -> DataType:
             return DataType(type='PositiveInt')
         if len(data_type_kwargs) == 1 and data_type_kwargs.get('ge') == 0:
             return DataType(type='NegativeInt')
-        return DataType(type='conint', is_func=True, kwargs=data_type_kwargs)
+        return DataType(
+            type='conint',
+            is_func=True,
+            kwargs=data_type_kwargs,
+            imports_=[IMPORT_CONINT],
+        )
     return type_map[types]
 
 
@@ -102,7 +112,12 @@ def get_data_float_type(types: Types, **kwargs: Any) -> DataType:
             return DataType(type='PositiveFloat')
         if len(data_type_kwargs) == 1 and data_type_kwargs.get('ge') == 0:
             return DataType(type='NegativeFloat')
-        return DataType(type='confloat', is_func=True, kwargs=data_type_kwargs)
+        return DataType(
+            type='confloat',
+            is_func=True,
+            kwargs=data_type_kwargs,
+            imports_=[IMPORT_CONFLOAT],
+        )
     return type_map[types]
 
 

diff --git a/datamodel_code_generator/parser/openapi.py b/datamodel_code_generator/parser/openapi.py
@@ -144,6 +144,9 @@ def parse_all_of(self, name: str, obj: JsonSchemaObject) -> List[DataType]:
             custom_template_dir=self.custom_template_dir,
             extra_template_data=self.extra_template_data,
         )
+        # add imports for the fields
+        for f in fields:
+            data_model_type.imports.extend(f.imports)
         self.append_result(data_model_type)
 
         return [self.data_type(type=name, ref=True, version_compatible=True)]

diff --git a/tests/data/allof.yaml b/tests/data/allof.yaml
@@ -131,6 +131,17 @@ components:
           properties:
             number:
               type: string
+    AllOfCombine:
+      allOf:
+        - $ref: "#/components/schemas/Pet"
+        - type: object
+          properties:
+            birthdate:
+              type: string
+              format: date
+            size:
+              type: integer
+              minimum: 1
     AnyOfCombine:
         allOf:
           - $ref: "#/components/schemas/Pet"
@@ -168,6 +179,9 @@ components:
           properties:
             age:
               type: string
+            birthdate:
+              type: string
+              format: date-time
     Error:
       required:
         - code

diff --git a/tests/data/anyof.yaml b/tests/data/anyof.yaml
@@ -153,6 +153,9 @@ components:
               properties:
                 name:
                   type: string
+                birthday:
+                  type: string
+                  format: date
     Error:
       required:
         - code

diff --git a/tests/model/pydantic/test_types.py b/tests/model/pydantic/test_types.py
@@ -1,5 +1,9 @@
 import pytest
-from datamodel_code_generator.imports import IMPORT_CONSTR
+from datamodel_code_generator.imports import (
+    IMPORT_CONFLOAT,
+    IMPORT_CONINT,
+    IMPORT_CONSTR,
+)
 from datamodel_code_generator.model.pydantic.types import (
     get_data_float_type,
     get_data_int_type,
@@ -16,27 +20,40 @@
         (
             Types.integer,
             {'maximum': 10},
-            DataType(type='conint', is_func=True, kwargs={'gt': 10}),
+            DataType(
+                type='conint', is_func=True, kwargs={'gt': 10}, imports_=[IMPORT_CONINT]
+            ),
         ),
         (
             Types.integer,
             {'exclusiveMaximum': 10},
-            DataType(type='conint', is_func=True, kwargs={'ge': 10}),
+            DataType(
+                type='conint', is_func=True, kwargs={'ge': 10}, imports_=[IMPORT_CONINT]
+            ),
         ),
         (
             Types.integer,
             {'minimum': 10},
-            DataType(type='conint', is_func=True, kwargs={'lt': 10}),
+            DataType(
+                type='conint', is_func=True, kwargs={'lt': 10}, imports_=[IMPORT_CONINT]
+            ),
         ),
         (
             Types.integer,
             {'exclusiveMinimum': 10},
-            DataType(type='conint', is_func=True, kwargs={'le': 10}),
+            DataType(
+                type='conint', is_func=True, kwargs={'le': 10}, imports_=[IMPORT_CONINT]
+            ),
         ),
         (
             Types.integer,
             {'multipleOf': 10},
-            DataType(type='conint', is_func=True, kwargs={'multiple_of': 10}),
+            DataType(
+                type='conint',
+                is_func=True,
+                kwargs={'multiple_of': 10},
+                imports_=[IMPORT_CONINT],
+            ),
         ),
         (Types.integer, {'exclusiveMinimum': 0}, DataType(type='PositiveInt')),
         (Types.integer, {'exclusiveMaximum': 0}, DataType(type='NegativeInt')),
@@ -53,27 +70,52 @@ def test_get_data_int_type(types, params, data_type):
         (
             Types.float,
             {'maximum': 10},
-            DataType(type='confloat', is_func=True, kwargs={'gt': 10}),
+            DataType(
+                type='confloat',
+                is_func=True,
+                kwargs={'gt': 10},
+                imports_=[IMPORT_CONFLOAT],
+            ),
         ),
         (
             Types.float,
             {'exclusiveMaximum': 10},
-            DataType(type='confloat', is_func=True, kwargs={'ge': 10}),
+            DataType(
+                type='confloat',
+                is_func=True,
+                kwargs={'ge': 10},
+                imports_=[IMPORT_CONFLOAT],
+            ),
         ),
         (
             Types.float,
             {'minimum': 10},
-            DataType(type='confloat', is_func=True, kwargs={'lt': 10}),
+            DataType(
+                type='confloat',
+                is_func=True,
+                kwargs={'lt': 10},
+                imports_=[IMPORT_CONFLOAT],
+            ),
         ),
         (
             Types.float,
             {'exclusiveMinimum': 10},
-            DataType(type='confloat', is_func=True, kwargs={'le': 10}),
+            DataType(
+                type='confloat',
+                is_func=True,
+                kwargs={'le': 10},
+                imports_=[IMPORT_CONFLOAT],
+            ),
         ),
         (
             Types.float,
             {'multipleOf': 10},
-            DataType(type='confloat', is_func=True, kwargs={'multiple_of': 10}),
+            DataType(
+                type='confloat',
+                is_func=True,
+                kwargs={'multiple_of': 10},
+                imports_=[IMPORT_CONFLOAT],
+            ),
         ),
         (Types.float, {'exclusiveMinimum': 0}, DataType(type='PositiveFloat')),
         (Types.float, {'exclusiveMaximum': 0}, DataType(type='NegativeFloat')),

diff --git a/tests/parser/test_openapi.py b/tests/parser/test_openapi.py
@@ -771,6 +771,7 @@ def test_openapi_parser_parse_anyof():
         parser.parse()
         == '''from __future__ import annotations
 
+from datetime import date
 from typing import List, Optional, Union
 
 from pydantic import BaseModel
@@ -806,6 +807,7 @@ class AnyOfobj(BaseModel):
 
 class AnyOfArrayItem(BaseModel):
     name: Optional[str] = None
+    birthday: Optional[date] = None
 
 
 class AnyOfArray(BaseModel):
@@ -827,9 +829,10 @@ def test_openapi_parser_parse_allof():
         parser.parse()
         == '''from __future__ import annotations
 
+from datetime import date, datetime
 from typing import List, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, conint
 
 
 class Pet(BaseModel):
@@ -851,6 +854,11 @@ class AllOfobj(BaseModel):
     number: Optional[str] = None
 
 
+class AllOfCombine(Pet):
+    birthdate: Optional[date] = None
+    size: Optional[conint(lt=1.0)] = None
+
+
 class AnyOfCombine(Pet, Car):
     age: Optional[str] = None
 
@@ -873,6 +881,7 @@ class AnyOfCombineInArray(BaseModel):
 
 class AnyOfCombineInRoot(Pet, Car):
     age: Optional[str] = None
+    birthdate: Optional[datetime] = None
 
 
 class Error(BaseModel):