support $id (#260)

* support $id

* refactor code

* refactor code

README.md

@@ -46,7 +46,7 @@ See [documentation](https://koxudaxi.github.io/datamodel-code-generator) for mor
 - anyOf (as Union)
 - oneOf (as Union)
 - $ref ([http extra](#http-extra-option) is required when resolving $ref for remote files.)
-
+- $id (for [JSONSchema](https://json-schema.org/understanding-json-schema/structuring.html#the-id-property))
 
 ## Installation
 

datamodel_code_generator/parser/base.py

@@ -320,11 +320,17 @@ def parse(
                             reference = self.model_resolver.get(
                                 data_type.reference.path
                             )
-                            if (
-                                reference
-                                and reference.actual_module_name == module_path
+                            if reference and (
+                                (
+                                    isinstance(self.source, Path)
+                                    and self.source.is_file()
+                                    and self.source.name
+                                    == reference.path.split('#/')[0]
+                                )
+                                or reference.actual_module_name == module_path
                             ):
-                                model.reference_classes.remove(name)
+                                if name in model.reference_classes:  # pragma: no cover
+                                    model.reference_classes.remove(name)
                                 continue
                         if full_path in alias_map:
                             alias = alias_map[full_path] or import_

datamodel_code_generator/parser/jsonschema.py

@@ -686,7 +686,7 @@ def _get_ref_body(self, ref: str) -> Dict[Any, Any]:
                 ref_body = yaml.safe_load(raw_body)
                 self.remote_object_cache[ref] = ref_body
         else:
-            if remote_object:  # pragma: no cover
+            if remote_object:
                 ref_body = remote_object
             else:
                 # Remote Reference – $ref: 'document.json' Uses the whole document located on the same server and in
@@ -751,6 +751,25 @@ def parse_ref(self, obj: JsonSchemaObject, path: List[str]) -> None:
             for value in obj.properties.values():
                 self.parse_ref(value, path)
 
+    def parse_id(self, obj: JsonSchemaObject, path: List[str]) -> None:
+        if obj.id:
+            self.model_resolver.add_id(obj.id, path)
+        if obj.items:
+            if isinstance(obj.items, JsonSchemaObject):
+                self.parse_id(obj.items, path)
+            else:
+                for item in obj.items:
+                    self.parse_id(item, path)
+        if isinstance(obj.additionalProperties, JsonSchemaObject):  # pragma: no cover
+            self.parse_id(obj.additionalProperties, path)
+        for item in obj.anyOf:  # pragma: no cover
+            self.parse_id(item, path)  # pragma: no cover
+        for item in obj.allOf:
+            self.parse_id(item, path)  # pragma: no cover
+        if obj.properties:
+            for value in obj.properties.values():
+                self.parse_id(value, path)
+
     @contextmanager
     def root_id_context(self, root_raw: Dict[str, Any]) -> Generator[None, None, None]:
         root_id: Optional[str] = root_raw.get('$id')
@@ -769,8 +788,9 @@ def root_id_context(self, root_raw: Dict[str, Any]) -> Generator[None, None, Non
         self.root_id = previous_root_id
 
     def parse_raw_obj(self, name: str, raw: Dict[str, Any], path: List[str],) -> None:
-        obj = JsonSchemaObject.parse_obj(raw)
+        self.parse_obj(name, JsonSchemaObject.parse_obj(raw), path)
 
+    def parse_obj(self, name: str, obj: JsonSchemaObject, path: List[str],) -> None:
         name = self.model_resolver.add(path, name, class_name=True).name
         if obj.is_object:
             self.parse_object(name, obj, path)
@@ -804,8 +824,17 @@ def parse_raw(self) -> None:
                     raise InvalidClassNameError(obj_name)
 
             obj_name = self.model_resolver.add(path_parts, obj_name, unique=False).name
+            root_obj = JsonSchemaObject.parse_obj(self.raw_obj)
             with self.root_id_context(self.raw_obj):
-                self.parse_raw_obj(obj_name, self.raw_obj, path_parts)
+                definitions = self.raw_obj.get('definitions', {})
+
+                # parse $id before parsing $ref
+                self.parse_id(root_obj, path_parts)
+                for key, model in definitions.items():
+                    obj = JsonSchemaObject.parse_obj(model)
+                    self.parse_id(obj, [*path_parts, '#/definitions', key])
+
+                self.parse_obj(obj_name, root_obj, path_parts)
                 definitions = self.raw_obj.get('definitions', {})
                 for key, model in definitions.items():
                     self.parse_raw_obj(key, model, [*path_parts, '#/definitions', key])

datamodel_code_generator/reference.py

@@ -1,7 +1,8 @@
 import re
+from collections import defaultdict
 from keyword import iskeyword
 from pathlib import Path
-from typing import Dict, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import DefaultDict, Dict, List, Mapping, Optional, Pattern, Tuple, Union
 
 import inflect
 from pydantic import BaseModel
@@ -26,29 +27,38 @@ def module_name(self) -> Optional[str]:
         return module_name
 
 
+ID_PATTERN: Pattern[str] = re.compile(r'^#[^/].*')
+
+
 class ModelResolver:
     def __init__(self, aliases: Optional[Mapping[str, str]] = None) -> None:
         self.references: Dict[str, Reference] = {}
         self.aliases: Mapping[str, str] = {**aliases} if aliases is not None else {}
         self._current_root: List[str] = []
         self._root_id_base_path: Optional[str] = None
+        self.ids: DefaultDict[str, Dict[str, str]] = defaultdict(dict)
 
     @property
     def current_root(self) -> List[str]:
-        return self._current_root  # pragma: no cover
+        return self._current_root
 
     def set_current_root(self, current_root: List[str]) -> None:
         self._current_root = current_root
 
     @property
     def root_id_base_path(self) -> Optional[str]:
-        return self._root_id_base_path  # pragma: no cover
+        return self._root_id_base_path
 
     def set_root_id_base_path(self, root_id_base_path: Optional[str]) -> None:
         self._root_id_base_path = root_id_base_path
 
+    def add_id(self, id_: str, path: List[str]) -> None:
+        self.ids['/'.join(self.current_root)][id_] = self._get_path(path)
+
     def _get_path(self, path: List[str]) -> str:
         joined_path = '/'.join(p for p in path if p).replace('/#', '#')
+        if ID_PATTERN.match(joined_path):
+            return self.ids['/'.join(self.current_root)][joined_path]
         if '#' in joined_path:
             delimiter = joined_path.index('#')
             return f"{''.join(joined_path[:delimiter])}#{''.join(joined_path[delimiter + 1:])}"
@@ -65,13 +75,7 @@ def add_ref(self, ref: str, actual_module_name: Optional[str] = None) -> Referen
             return reference
         split_ref = ref.rsplit('/', 1)
         if len(split_ref) == 1:
-            first_ref = split_ref[0]
-            if first_ref[0] == '#':
-                # TODO Support $id with $ref
-                # https://json-schema.org/understanding-json-schema/structuring.html#using-id-with-ref
-                raise NotImplementedError('It is not support to use $id with $ref')
-            else:
-                parents, original_name = self.root_id_base_path, Path(first_ref).stem
+            parents, original_name = self.root_id_base_path, Path(split_ref[0]).stem
         else:
             parents, original_name = split_ref
         loaded: bool = not ref.startswith(('https://', 'http://'))
@@ -106,7 +110,7 @@ def add(
         if not original_name:
             original_name = Path(joined_path.split('#')[0]).stem
         name = original_name
-        if singular_name:  # pragma: no cover
+        if singular_name:
             name = get_singular_name(name, singular_name_suffix)
         if class_name:
             name = self.get_class_name(name, unique)
@@ -156,6 +160,8 @@ def validate_name(cls, name: str) -> bool:
         return name.isidentifier() and not iskeyword(name)
 
     def get_valid_name(self, name: str, camel: bool = False) -> str:
+        if name[0] == '#':
+            name = name[1:]
         # TODO: when first character is a number
         replaced_name = re.sub(r'\W', '_', name)
         if re.match(r'^[0-9]', replaced_name):

docs/support-data-types.md

@@ -28,3 +28,4 @@ This codegen supports major data types to OpenAPI/JSON Schema
 - anyOf (as Union)
 - oneOf (as Union)
 - $ref ([http extra](../#http-extra-option) is required when resolving $ref for remote files.)
+- $id (for [JSONSchema](https://json-schema.org/understanding-json-schema/structuring.html#the-id-property))

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

@@ -0,0 +1,20 @@
+# generated by datamodel-codegen:
+#   filename:  id.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class Address(BaseModel):
+    street_address: str
+    city: str
+    state: str
+
+
+class Model(BaseModel):
+    billing_address: Optional[Address] = None
+    shipping_address: Optional[Address] = None

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

@@ -0,0 +1,20 @@
+# generated by datamodel-codegen:
+#   filename:  <stdin>
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class Address(BaseModel):
+    street_address: str
+    city: str
+    state: str
+
+
+class Model(BaseModel):
+    billing_address: Optional[Address] = None
+    shipping_address: Optional[Address] = None

tests/data/jsonschema/id.json

@@ -0,0 +1,23 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+
+  "definitions": {
+    "address": {
+      "$id": "#address",
+      "type": "object",
+      "properties": {
+        "street_address": { "type": "string" },
+        "city":           { "type": "string" },
+        "state":          { "type": "string" }
+      },
+      "required": ["street_address", "city", "state"]
+    }
+  },
+
+  "type": "object",
+
+  "properties": {
+    "billing_address": { "$ref": "#address" },
+    "shipping_address": { "$ref": "#address" }
+  }
+}

tests/test_main.py

@@ -996,3 +996,45 @@ def test_main_root_id_jsonschema_root_id_failed(mocker):
         )
     with pytest.raises(SystemExit):
         main()
+
+
+@freeze_time('2019-07-26')
+def test_main_jsonschema_id():
+    with TemporaryDirectory() as output_dir:
+        output_file: Path = Path(output_dir) / 'output.py'
+        return_code: Exit = main(
+            [
+                '--input',
+                str(JSON_SCHEMA_DATA_PATH / 'id.json'),
+                '--output',
+                str(output_file),
+                '--input-file-type',
+                'jsonschema',
+            ]
+        )
+        assert return_code == Exit.OK
+        assert (
+            output_file.read_text()
+            == (EXPECTED_MAIN_PATH / 'main_jsonschema_id' / 'output.py').read_text()
+        )
+    with pytest.raises(SystemExit):
+        main()
+
+
+@freeze_time('2019-07-26')
+def test_main_jsonschema_id_as_stdin(monkeypatch):
+    with TemporaryDirectory() as output_dir:
+        output_file: Path = Path(output_dir) / 'output.py'
+        monkeypatch.setattr('sys.stdin', (JSON_SCHEMA_DATA_PATH / 'id.json').open())
+        return_code: Exit = main(
+            ['--output', str(output_file), '--input-file-type', 'jsonschema',]
+        )
+        assert return_code == Exit.OK
+        assert (
+            output_file.read_text()
+            == (
+                EXPECTED_MAIN_PATH / 'main_jsonschema_id_stdin' / 'output.py'
+            ).read_text()
+        )
+    with pytest.raises(SystemExit):
+        main()