improve reusing model (#297)

datamodel_code_generator/parser/base.py

@@ -27,6 +27,7 @@
 from ..imports import IMPORT_ANNOTATIONS, Import, Imports
 from ..model import pydantic as pydantic_model
 from ..model.base import ALL_MODEL, DataModel, DataModelFieldBase
+from ..model.enum import Enum
 from ..reference import ModelResolver
 from ..types import DataType, DataTypeManager
 
@@ -446,7 +447,19 @@ def parse(
                     )
                     cached_model_name = model_cache.get(model_key)
                     if cached_model_name:
-                        duplicated_model_names[model.name] = cached_model_name
+                        if isinstance(model, Enum):
+                            duplicated_model_names[model.name] = cached_model_name
+                        else:
+                            index = models.index(model)
+                            inherited_model = model.__class__(
+                                name=model.name,
+                                fields=[],
+                                base_classes=[cached_model_name],
+                                description=model.description,
+                            )
+                            models.insert(index, inherited_model)
+                            models.remove(model)
+
                     else:
                         model_cache[model_key] = model.name
 

tests/data/expected/main/main_json_reuse_enum/output.py

@@ -0,0 +1,28 @@
+# generated by datamodel-codegen:
+#   filename:  duplicate_enum.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class Animal(Enum):
+    dog = 'dog'
+    cat = 'cat'
+    snake = 'snake'
+
+
+class Pet(Enum):
+    dog = 'dog'
+    cat = 'cat'
+    snake = 'snake'
+
+
+class User(BaseModel):
+    name: Optional[str] = None
+    animal: Optional[Animal] = None
+    pet: Optional[Animal] = None

tests/data/expected/main/main_json_resuse_model/output.py → tests/data/expected/main/main_json_reuse_model/output.py

@@ -13,10 +13,8 @@ class ArmRight(BaseModel):
     Joint_3: int = Field(..., alias='Joint 3')
 
 
-class ArmLeft(BaseModel):
-    Joint_1: int = Field(..., alias='Joint 1')
-    Joint_2: int = Field(..., alias='Joint 2')
-    Joint_3: int = Field(..., alias='Joint 3')
+class ArmLeft(ArmRight):
+    pass
 
 
 class Head(BaseModel):
@@ -25,5 +23,5 @@ class Head(BaseModel):
 
 class Model(BaseModel):
     Arm_Right: ArmRight = Field(..., alias='Arm Right')
-    Arm_Left: ArmRight = Field(..., alias='Arm Left')
+    Arm_Left: ArmLeft = Field(..., alias='Arm Left')
     Head: Head

tests/data/jsonschema/duplicate_enum.json

@@ -0,0 +1,26 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema",
+  "title": "User",
+  "type": "object",
+  "properties": {
+      "name": {
+          "type": "string"
+      },
+      "animal": {
+        "type": "string",
+        "enum": [
+            "dog",
+            "cat",
+            "snake"
+        ]
+    },
+      "pet": {
+        "type": "string",
+        "enum": [
+            "dog",
+            "cat",
+            "snake"
+        ]
+    }
+  }
+}

tests/test_main.py

@@ -1269,7 +1269,31 @@ def test_main_json_reuse_model():
         assert return_code == Exit.OK
         assert (
             output_file.read_text()
-            == (EXPECTED_MAIN_PATH / 'main_json_resuse_model' / 'output.py').read_text()
+            == (EXPECTED_MAIN_PATH / 'main_json_reuse_model' / 'output.py').read_text()
+        )
+    with pytest.raises(SystemExit):
+        main()
+
+
+@freeze_time('2019-07-26')
+def test_main_json_reuse_enum():
+    with TemporaryDirectory() as output_dir:
+        output_file: Path = Path(output_dir) / 'output.py'
+        return_code: Exit = main(
+            [
+                '--input',
+                str(JSON_SCHEMA_DATA_PATH / 'duplicate_enum.json'),
+                '--output',
+                str(output_file),
+                '--input-file-type',
+                'jsonschema',
+                '--reuse-model',
+            ]
+        )
+        assert return_code == Exit.OK
+        assert (
+            output_file.read_text()
+            == (EXPECTED_MAIN_PATH / 'main_json_reuse_enum' / 'output.py').read_text()
         )
     with pytest.raises(SystemExit):
         main()