Support $defs for Draft 2019-09 (#1630)

* Support $defs * Add unittest * Add unittest * Add unittest
koxudaxi · Oct 22, 2023 · 39c46f9 · 39c46f9
1 parent 405b50a
commit 39c46f9
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 9 deletions.
diff --git a/datamodel_code_generator/parser/jsonschema.py b/datamodel_code_generator/parser/jsonschema.py
@@ -356,7 +356,7 @@ def _get_type(type_: str, format__: Optional[str] = None) -> Types:
 
 @snooper_to_methods(max_variable_length=None)
 class JsonSchemaParser(Parser):
-    SCHEMA_PATH: ClassVar[str] = '#/definitions'
+    SCHEMA_PATHS: ClassVar[List[str]] = ['#/definitions', '#/$defs']
 
     def __init__(
         self,
@@ -516,8 +516,8 @@ def get_field_extras(self, obj: JsonSchemaObject) -> Dict[str, Any]:
             }
 
     @cached_property
-    def schema_paths(self) -> List[str]:
-        return self.SCHEMA_PATH.lstrip('#/').split('/')
+    def schema_paths(self) -> List[Tuple[str, List[str]]]:
+        return [(s, s.lstrip('#/').split('/')) for s in self.SCHEMA_PATHS]
 
     @property
     def root_id(self) -> Optional[str]:
@@ -1664,13 +1664,20 @@ def _parse_file(
                 # parse $id before parsing $ref
                 root_obj = JsonSchemaObject.parse_obj(raw)
                 self.parse_id(root_obj, path_parts)
-                try:
-                    definitions = get_model_by_path(raw, self.schema_paths)
-                except KeyError:
+                definitions: Optional[Dict[Any, Any]] = None
+                for schema_path, split_schema_path in self.schema_paths:
+                    try:
+                        definitions = get_model_by_path(raw, split_schema_path)
+                        if definitions:
+                            break
+                    except KeyError:
+                        continue
+                if definitions is None:
                     definitions = {}
+
                 for key, model in definitions.items():
                     obj = JsonSchemaObject.parse_obj(model)
-                    self.parse_id(obj, [*path_parts, self.SCHEMA_PATH, key])
+                    self.parse_id(obj, [*path_parts, schema_path, key])
 
                 if object_paths:
                     models = get_model_by_path(raw, object_paths)
@@ -1679,7 +1686,7 @@ def _parse_file(
                 else:
                     self.parse_obj(obj_name, root_obj, path_parts or ['#'])
                 for key, model in definitions.items():
-                    path = [*path_parts, self.SCHEMA_PATH, key]
+                    path = [*path_parts, schema_path, key]
                     reference = self.model_resolver.get(path)
                     if not reference or not reference.loaded:
                         self.parse_raw_obj(key, model, path)

diff --git a/datamodel_code_generator/parser/openapi.py b/datamodel_code_generator/parser/openapi.py
@@ -150,7 +150,7 @@ class ComponentsObject(BaseModel):
 
 @snooper_to_methods(max_variable_length=None)
 class OpenAPIParser(JsonSchemaParser):
-    SCHEMA_PATH: ClassVar[str] = '#/components/schemas'
+    SCHEMA_PATHS: ClassVar[List[str]] = ['#/components/schemas']
 
     def __init__(
         self,

diff --git a/tests/data/jsonschema/user_defs.json b/tests/data/jsonschema/user_defs.json
@@ -0,0 +1,33 @@
+{
+  "$schema": "https://json-schema.org/draft/2019-09/schema",
+  "$defs": {
+    "User": {
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string",
+          "example": "ken"
+        },
+        "pets": {
+          "type": "array",
+          "items": {
+            "$ref": "#/$defs/User"
+          },
+          "default_factory": "list"
+        }
+      }
+    },
+    "Pet": {
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string",
+          "examples": [
+            "dog",
+            "cat"
+          ]
+        }
+      }
+    }
+  }
+}
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -5879,3 +5879,23 @@ def test_main_duplicate_field_constraints_msgspec(
                 path.relative_to(main_modular_dir)
             ).read_text()
             assert result == path.read_text()
+
+
+@freeze_time('2019-07-26')
+def test_main_dataclass_field_defs():
+    with TemporaryDirectory() as output_dir:
+        output_file: Path = Path(output_dir) / 'output.py'
+        return_code: Exit = main(
+            [
+                '--input',
+                str(JSON_SCHEMA_DATA_PATH / 'user_defs.json'),
+                '--output',
+                str(output_file),
+                '--output-model-type',
+                'dataclasses.dataclass',
+            ]
+        )
+        assert return_code == Exit.OK
+        assert output_file.read_text() == (
+            EXPECTED_MAIN_PATH / 'main_dataclass_field' / 'output.py'
+        ).read_text().replace('filename:  user.json', 'filename:  user_defs.json')