fix multiple sub-type in List (#150)

koxudaxi · Jun 12, 2020 · 26cc2b4 · 26cc2b4
1 parent abde9b4
commit 26cc2b4
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 7 deletions.
diff --git a/datamodel_code_generator/model/base.py b/datamodel_code_generator/model/base.py
@@ -74,10 +74,8 @@ def _get_type_hint(self) -> Optional[str]:
             return type_hint
         if self.is_list:
             self.imports.append(IMPORT_LIST)
-            if self.is_union:
-                self.imports.append(IMPORT_UNION)
-                return f'{LIST}[{UNION}[{type_hint}]]'
-            return f'{LIST}[{type_hint}]'
+            self.imports.append(IMPORT_UNION)
+            return f'{LIST}[{UNION}[{type_hint}]]'
         self.imports.append(IMPORT_UNION)
         return f'{UNION}[{type_hint}]'
 

diff --git a/datamodel_code_generator/parser/jsonschema.py b/datamodel_code_generator/parser/jsonschema.py
@@ -221,9 +221,12 @@ def parse_any_of(self, name: str, obj: JsonSchemaObject) -> List[DataType]:
                 )
             ):
                 # trivial item types
+                types = [t.type_hint for t in self.get_data_type(any_of_item.items)]
                 any_of_data_types.append(
                     self.data_type(
-                        type=f"List[{', '.join([t.type_hint for t in self.get_data_type(any_of_item.items)])}]",
+                        type=f"List[Union[{', '.join(types)}]]"
+                        if len(types) > 1
+                        else f"List[{types[0]}]",
                         imports_=[Import(from_='typing', import_='List')],
                     )
                 )

diff --git a/tests/model/test_base.py b/tests/model/test_base.py
@@ -239,7 +239,7 @@ def test_data_field():
         is_list=True,
         is_union=False,
     )
-    assert field.type_hint == 'List[str, int]'
+    assert field.type_hint == 'List[Union[str, int]]'
     field = DataModelFieldBase(
         name='a',
         data_types=[DataType(type='str'), DataType(type='int')],
@@ -277,7 +277,7 @@ def test_data_field():
         is_list=True,
         is_union=False,
     )
-    assert field.type_hint == 'Optional[List[str, int]]'
+    assert field.type_hint == 'Optional[List[Union[str, int]]]'
 
     field = DataModelFieldBase(
         name='a', data_types=[], required=False, is_list=True, is_union=False