Skip to content

Commit 8c7550c

Browse files
authored
Expose schema $id and path to template context (#2798)
* Expose schema $id and path to template context * Replace unit tests with e2e test for schema_id feature * Remove duplicate set_schema_id call * Update test template with escape_docstring filter * Use simple class attributes instead of ClassVar in test template * Refactor _set_schema_metadata to call individual setter methods
1 parent 5aeb0b6 commit 8c7550c

6 files changed

Lines changed: 150 additions & 7 deletions

File tree

src/datamodel_code_generator/model/base.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -887,6 +887,7 @@ def render(self, *, class_name: str | None = None) -> str:
887887
methods=self.methods,
888888
description=self.description,
889889
dataclass_arguments=self.dataclass_arguments,
890+
path=self.path,
890891
**self.extra_template_data,
891892
)
892893

src/datamodel_code_generator/parser/jsonschema.py

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1168,14 +1168,17 @@ def set_title(self, path: str, obj: JsonSchemaObject) -> None:
11681168
if obj.title:
11691169
self.extra_template_data[path]["title"] = obj.title
11701170

1171+
def set_schema_id(self, path: str, obj: JsonSchemaObject) -> None:
1172+
"""Set $id in extra template data."""
1173+
if obj.id:
1174+
self.extra_template_data[path]["schema_id"] = obj.id
1175+
11711176
def _set_schema_metadata(self, path: str, obj: JsonSchemaObject) -> None:
1172-
"""Set title, additionalProperties and unevaluatedProperties in extra template data."""
1173-
if obj.title:
1174-
self.extra_template_data[path]["title"] = obj.title
1175-
if isinstance(obj.additionalProperties, bool):
1176-
self.extra_template_data[path]["additionalProperties"] = obj.additionalProperties
1177-
if isinstance(obj.unevaluatedProperties, bool):
1178-
self.extra_template_data[path]["unevaluatedProperties"] = obj.unevaluatedProperties
1177+
"""Set title, $id, additionalProperties and unevaluatedProperties in extra template data."""
1178+
self.set_title(path, obj)
1179+
self.set_schema_id(path, obj)
1180+
self.set_additional_properties(path, obj)
1181+
self.set_unevaluated_properties(path, obj)
11791182

11801183
def set_schema_extensions(self, path: str, obj: JsonSchemaObject) -> None:
11811184
"""Set schema extensions (x-* fields) in extra template data."""
@@ -1926,6 +1929,7 @@ def _parse_object_common_part( # noqa: PLR0912, PLR0913, PLR0915
19261929
reference = self.model_resolver.add(path, name, class_name=True, loaded=True)
19271930
self.set_additional_properties(reference.path, obj)
19281931
self.set_unevaluated_properties(reference.path, obj)
1932+
self.set_schema_id(reference.path, obj)
19291933
self.set_schema_extensions(reference.path, obj)
19301934

19311935
generates_separate = self._should_generate_separate_models(fields, base_classes)
@@ -2268,6 +2272,7 @@ def parse_object(
22682272
)
22692273
class_name = reference.name
22702274
self.set_title(reference.path, obj)
2275+
self.set_schema_id(reference.path, obj)
22712276
if self.read_only_write_only_model_type is not None and obj.properties:
22722277
for prop in obj.properties.values():
22732278
if isinstance(prop, JsonSchemaObject) and prop.ref:
Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
# generated by datamodel-codegen:
2+
# filename: schema_id.json
3+
# timestamp: 1985-10-26T08:21:00+00:00
4+
5+
from __future__ import annotations
6+
7+
from pydantic import BaseModel
8+
9+
10+
class Model(BaseModel):
11+
__schema_id__ = "https://example.com/schemas/user"
12+
__path__ = "#"
13+
id: str
14+
name: str
15+
16+
17+
class Address(BaseModel):
18+
__schema_id__ = "#address"
19+
__path__ = "#/definitions/Address"
20+
street: str
21+
city: str | None = None
Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
{
2+
"$schema": "http://json-schema.org/draft-07/schema#",
3+
"$id": "https://example.com/schemas/user",
4+
"type": "object",
5+
"properties": {
6+
"id": {
7+
"type": "string"
8+
},
9+
"name": {
10+
"type": "string"
11+
}
12+
},
13+
"required": ["id", "name"],
14+
"definitions": {
15+
"Address": {
16+
"$id": "#address",
17+
"type": "object",
18+
"properties": {
19+
"street": {
20+
"type": "string"
21+
},
22+
"city": {
23+
"type": "string"
24+
}
25+
},
26+
"required": ["street"]
27+
}
28+
}
29+
}
Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
{% if base_class != "BaseModel" and "," not in base_class and not fields and not config and not description -%}
2+
3+
{# if this is just going to be `class Foo(Bar): pass`, then might as well just make Foo
4+
an alias for Bar: every pydantic model class consumes considerable memory. #}
5+
{{ class_name }} = {{ base_class }}
6+
7+
{% else -%}
8+
9+
{% for decorator in decorators -%}
10+
{{ decorator }}
11+
{% endfor -%}
12+
class {{ class_name }}({{ base_class }}):{% if comment is defined %} # {{ comment }}{% endif %}
13+
{%- if description %}
14+
"""
15+
{{ description | escape_docstring | indent(4) }}
16+
"""
17+
{%- endif %}
18+
{%- if schema_id is defined and schema_id %}
19+
__schema_id__ = "{{ schema_id }}"
20+
{%- endif %}
21+
{%- if path is defined and path %}
22+
__path__ = "{{ path }}"
23+
{%- endif %}
24+
{%- if not fields and not description and not config and not (schema_id is defined and schema_id) %}
25+
pass
26+
{%- endif %}
27+
{%- if config %}
28+
{%- filter indent(4) %}
29+
{% include 'ConfigDict.jinja2' %}
30+
{%- endfilter %}
31+
{%- endif %}
32+
{%- for field in fields %}
33+
{%- if not field.annotated and field.field %}
34+
{{ field.name }}: {{ field.type_hint }} = {{ field.field }}
35+
{%- else %}
36+
{%- if field.annotated %}
37+
{{ field.name }}: {{ field.annotated }}
38+
{%- else %}
39+
{{ field.name }}: {{ field.type_hint }}
40+
{%- endif %}
41+
{%- if not field.has_default_factory_in_field and not field.required and (field.represented_default != 'None' or not field.strip_default_none or field.data_type.is_optional)
42+
%} = {{ field.represented_default }}
43+
{%- endif -%}
44+
{%- endif %}
45+
{%- if field.docstring %}
46+
"""
47+
{{ field.docstring | escape_docstring | indent(4) }}
48+
"""
49+
{%- if field.use_inline_field_description and not loop.last %}
50+
51+
{% endif %}
52+
{%- elif field.inline_field_docstring %}
53+
{{ field.inline_field_docstring }}
54+
{%- if not loop.last %}
55+
56+
{% endif %}
57+
{%- endif %}
58+
{%- for method in methods -%}
59+
{{ method }}
60+
{%- endfor -%}
61+
{%- endfor -%}
62+
63+
{%- endif %}

tests/main/jsonschema/test_main_jsonschema.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
)
2424
from datamodel_code_generator.__main__ import Exit, main
2525
from datamodel_code_generator.format import is_supported_in_black
26+
from datamodel_code_generator.model import base as model_base
2627
from tests.conftest import assert_directory_content, freeze_time
2728
from tests.main.conftest import (
2829
ALIASES_DATA_PATH,
@@ -6342,3 +6343,26 @@ def test_main_use_root_model_type_alias(output_file: Path) -> None:
63426343
"3.10",
63436344
],
63446345
)
6346+
6347+
6348+
def test_main_jsonschema_schema_id(
6349+
capsys: pytest.CaptureFixture, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
6350+
) -> None:
6351+
"""Test that $id is exposed as schema_id in custom templates (issue #2098)."""
6352+
model_base._get_environment.cache_clear()
6353+
model_base._get_template_with_custom_dir.cache_clear()
6354+
monkeypatch.chdir(tmp_path)
6355+
with freeze_time(TIMESTAMP):
6356+
run_main_and_assert(
6357+
input_path=JSON_SCHEMA_DATA_PATH / "schema_id.json",
6358+
output_path=None,
6359+
expected_stdout_path=EXPECTED_JSON_SCHEMA_PATH / "schema_id.py",
6360+
capsys=capsys,
6361+
input_file_type=None,
6362+
extra_args=[
6363+
"--custom-template-dir",
6364+
str(DATA_PATH / "templates_schema_id"),
6365+
"--output-model-type",
6366+
"pydantic_v2.BaseModel",
6367+
],
6368+
)

0 commit comments

Comments
 (0)