diff --git a/src/datamodel_code_generator/parser/jsonschema.py b/src/datamodel_code_generator/parser/jsonschema.py index 37f7828e0..d50154e18 100644 --- a/src/datamodel_code_generator/parser/jsonschema.py +++ b/src/datamodel_code_generator/parser/jsonschema.py @@ -46,7 +46,7 @@ PythonVersion, PythonVersionMin, ) -from datamodel_code_generator.imports import IMPORT_ANY +from datamodel_code_generator.imports import IMPORT_ANY, Import from datamodel_code_generator.model import DataModel, DataModelFieldBase from datamodel_code_generator.model import pydantic as pydantic_model from datamodel_code_generator.model.base import UNDEFINED, get_module_name, sanitize_module_name @@ -534,6 +534,42 @@ class JsonSchemaParser(Parser): SCHEMA_PATHS: ClassVar[list[str]] = ["#/definitions", "#/$defs"] SCHEMA_OBJECT_TYPE: ClassVar[type[JsonSchemaObject]] = JsonSchemaObject + COMPATIBLE_PYTHON_TYPES: ClassVar[dict[str, frozenset[str]]] = { + "string": frozenset({"str", "String"}), + "integer": frozenset({"int", "Integer"}), + "number": frozenset({"float", "int", "Number"}), + "boolean": frozenset({"bool", "Boolean"}), + "array": frozenset({ + "list", + "List", + "set", + "Set", + "frozenset", + "FrozenSet", + "Sequence", + "MutableSequence", + "tuple", + "Tuple", + "AbstractSet", + "MutableSet", + }), + "object": frozenset({"dict", "Dict", "Mapping", "MutableMapping", "TypedDict"}), + } + + PYTHON_TYPE_IMPORTS: ClassVar[dict[str, Import]] = { + "Callable": Import.from_full_path("collections.abc.Callable"), + "Iterable": Import.from_full_path("collections.abc.Iterable"), + "Iterator": Import.from_full_path("collections.abc.Iterator"), + "Generator": Import.from_full_path("collections.abc.Generator"), + "Awaitable": Import.from_full_path("collections.abc.Awaitable"), + "Coroutine": Import.from_full_path("collections.abc.Coroutine"), + "AsyncIterable": Import.from_full_path("collections.abc.AsyncIterable"), + "AsyncIterator": Import.from_full_path("collections.abc.AsyncIterator"), + "AsyncGenerator": Import.from_full_path("collections.abc.AsyncGenerator"), + "Pattern": Import.from_full_path("re.Pattern"), + "Match": Import.from_full_path("re.Match"), + } + def __init__( # noqa: PLR0913 self, source: str | Path | list[Path] | ParseResult, @@ -1147,6 +1183,10 @@ def get_object_field( # noqa: PLR0913 def get_data_type(self, obj: JsonSchemaObject) -> DataType: """Get the data type for a JSON Schema object.""" + python_type_override = self._get_python_type_override(obj) + if python_type_override: + return python_type_override + if obj.type is None: if "const" in obj.extras: return self.data_type_manager.get_data_type_from_value(obj.extras["const"]) @@ -1276,6 +1316,49 @@ class decorator which does not preserve staticmethod descriptors. return {} + def _get_python_type_base(self, python_type: str) -> str: # noqa: PLR6301 + """Extract base type from a Python type annotation string.""" + if "." in python_type.split("[", maxsplit=1)[0]: + base = python_type.split("[", maxsplit=1)[0].rsplit(".", 1)[-1] + else: + base = python_type.split("[", maxsplit=1)[0].strip() + return base + + def _is_compatible_python_type(self, schema_type: str | None, python_type: str) -> bool: + """Check if x-python-type is compatible with the JSON Schema type.""" + if schema_type is None: + return True + base_type = self._get_python_type_base(python_type) + if base_type in {"Union", "Optional"}: + return True + compatible = self.COMPATIBLE_PYTHON_TYPES.get(schema_type, frozenset()) + return base_type in compatible + + def _get_python_type_override(self, obj: JsonSchemaObject) -> DataType | None: + """Get DataType from x-python-type if it's incompatible with schema type.""" + x_python_type = obj.extras.get("x-python-type") + if not x_python_type or not isinstance(x_python_type, str): + return None + + schema_type = obj.type if isinstance(obj.type, str) else None + if self._is_compatible_python_type(schema_type, x_python_type): + return None + + base_type = self._get_python_type_base(x_python_type) + import_ = self.PYTHON_TYPE_IMPORTS.get(base_type) + + # Convert fully qualified path to short name when import is added + type_str = x_python_type + prefix = x_python_type.split("[", maxsplit=1)[0] + if "." in prefix: + # Replace the fully qualified prefix with just the base type name + type_str = base_type + x_python_type[len(prefix) :] + if not import_: + # If not in predefined imports, create import from the full path + import_ = Import.from_full_path(prefix) + + return self.data_type(type=type_str, import_=import_) + def _apply_title_as_name(self, name: str, obj: JsonSchemaObject) -> str: """Apply title as name if use_title_as_name is enabled.""" if self.use_title_as_name and obj.title: diff --git a/tests/data/expected/main/jsonschema/x_python_type_callable.py b/tests/data/expected/main/jsonschema/x_python_type_callable.py new file mode 100644 index 000000000..761cf0223 --- /dev/null +++ b/tests/data/expected/main/jsonschema/x_python_type_callable.py @@ -0,0 +1,14 @@ +# generated by datamodel-codegen: +# filename: x_python_type_callable.json +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from collections.abc import Callable +from typing import TypedDict + +from typing_extensions import NotRequired + + +class Model(TypedDict): + callback: NotRequired[Callable[[str], str]] diff --git a/tests/data/expected/main/jsonschema/x_python_type_callable_anyof.py b/tests/data/expected/main/jsonschema/x_python_type_callable_anyof.py new file mode 100644 index 000000000..5be2c7405 --- /dev/null +++ b/tests/data/expected/main/jsonschema/x_python_type_callable_anyof.py @@ -0,0 +1,14 @@ +# generated by datamodel-codegen: +# filename: x_python_type_callable_anyof.json +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from collections.abc import Callable +from typing import TypedDict + +from typing_extensions import NotRequired + + +class Model(TypedDict): + callback: NotRequired[Callable[[str], str] | None] diff --git a/tests/data/expected/main/jsonschema/x_python_type_compatible_set.py b/tests/data/expected/main/jsonschema/x_python_type_compatible_set.py new file mode 100644 index 000000000..c4c2cf142 --- /dev/null +++ b/tests/data/expected/main/jsonschema/x_python_type_compatible_set.py @@ -0,0 +1,13 @@ +# generated by datamodel-codegen: +# filename: x_python_type_compatible_set.json +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from typing import TypedDict + +from typing_extensions import NotRequired + + +class Model(TypedDict): + items: NotRequired[set[str]] diff --git a/tests/data/expected/main/jsonschema/x_python_type_custom_fqpath.py b/tests/data/expected/main/jsonschema/x_python_type_custom_fqpath.py new file mode 100644 index 000000000..633ccdd92 --- /dev/null +++ b/tests/data/expected/main/jsonschema/x_python_type_custom_fqpath.py @@ -0,0 +1,14 @@ +# generated by datamodel-codegen: +# filename: x_python_type_custom_fqpath.json +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from typing import TypedDict + +from my.custom import Handler +from typing_extensions import NotRequired + + +class Model(TypedDict): + handler: NotRequired[Handler[[str], str]] diff --git a/tests/data/expected/main/jsonschema/x_python_type_fqpath.py b/tests/data/expected/main/jsonschema/x_python_type_fqpath.py new file mode 100644 index 000000000..b569be31e --- /dev/null +++ b/tests/data/expected/main/jsonschema/x_python_type_fqpath.py @@ -0,0 +1,14 @@ +# generated by datamodel-codegen: +# filename: x_python_type_fqpath.json +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from collections.abc import Callable +from typing import TypedDict + +from typing_extensions import NotRequired + + +class Model(TypedDict): + callback: NotRequired[Callable[[str], str]] diff --git a/tests/data/expected/main/jsonschema/x_python_type_no_schema_type.py b/tests/data/expected/main/jsonschema/x_python_type_no_schema_type.py new file mode 100644 index 000000000..7584ba5b8 --- /dev/null +++ b/tests/data/expected/main/jsonschema/x_python_type_no_schema_type.py @@ -0,0 +1,13 @@ +# generated by datamodel-codegen: +# filename: x_python_type_no_schema_type.json +# timestamp: 2019-07-26T00:00:00+00:00 + +from __future__ import annotations + +from typing import Any, TypedDict + +from typing_extensions import NotRequired + + +class Model(TypedDict): + callback: NotRequired[Any] diff --git a/tests/data/jsonschema/x_python_type_callable.json b/tests/data/jsonschema/x_python_type_callable.json new file mode 100644 index 000000000..1ca0561e0 --- /dev/null +++ b/tests/data/jsonschema/x_python_type_callable.json @@ -0,0 +1,9 @@ +{ + "type": "object", + "properties": { + "callback": { + "type": "string", + "x-python-type": "Callable[[str], str]" + } + } +} diff --git a/tests/data/jsonschema/x_python_type_callable_anyof.json b/tests/data/jsonschema/x_python_type_callable_anyof.json new file mode 100644 index 000000000..333585d97 --- /dev/null +++ b/tests/data/jsonschema/x_python_type_callable_anyof.json @@ -0,0 +1,11 @@ +{ + "type": "object", + "properties": { + "callback": { + "anyOf": [ + {"type": "string", "x-python-type": "Callable[[str], str]"}, + {"type": "null"} + ] + } + } +} diff --git a/tests/data/jsonschema/x_python_type_compatible_set.json b/tests/data/jsonschema/x_python_type_compatible_set.json new file mode 100644 index 000000000..fd53a5ebe --- /dev/null +++ b/tests/data/jsonschema/x_python_type_compatible_set.json @@ -0,0 +1,10 @@ +{ + "type": "object", + "properties": { + "items": { + "type": "array", + "items": {"type": "string"}, + "x-python-type": "Set[str]" + } + } +} diff --git a/tests/data/jsonschema/x_python_type_custom_fqpath.json b/tests/data/jsonschema/x_python_type_custom_fqpath.json new file mode 100644 index 000000000..74ed36e7c --- /dev/null +++ b/tests/data/jsonschema/x_python_type_custom_fqpath.json @@ -0,0 +1,9 @@ +{ + "type": "object", + "properties": { + "handler": { + "type": "string", + "x-python-type": "my.custom.Handler[[str], str]" + } + } +} diff --git a/tests/data/jsonschema/x_python_type_fqpath.json b/tests/data/jsonschema/x_python_type_fqpath.json new file mode 100644 index 000000000..1ab5b7442 --- /dev/null +++ b/tests/data/jsonschema/x_python_type_fqpath.json @@ -0,0 +1,9 @@ +{ + "type": "object", + "properties": { + "callback": { + "type": "string", + "x-python-type": "collections.abc.Callable[[str], str]" + } + } +} diff --git a/tests/data/jsonschema/x_python_type_no_schema_type.json b/tests/data/jsonschema/x_python_type_no_schema_type.json new file mode 100644 index 000000000..b082e9a73 --- /dev/null +++ b/tests/data/jsonschema/x_python_type_no_schema_type.json @@ -0,0 +1,8 @@ +{ + "type": "object", + "properties": { + "callback": { + "x-python-type": "Callable[[str], str]" + } + } +} diff --git a/tests/main/jsonschema/test_main_jsonschema.py b/tests/main/jsonschema/test_main_jsonschema.py index 104849688..63f47286b 100644 --- a/tests/main/jsonschema/test_main_jsonschema.py +++ b/tests/main/jsonschema/test_main_jsonschema.py @@ -6965,3 +6965,69 @@ def test_main_jsonschema_ref_to_json_list_file() -> None: input_=JSON_SCHEMA_DATA_PATH / "ref_to_json_list" / "main.json", input_file_type=InputFileType.JsonSchema, ) + + +def test_x_python_type_callable(output_file: Path) -> None: + """Test x-python-type with Callable preserves the Callable type.""" + run_main_and_assert( + input_path=JSON_SCHEMA_DATA_PATH / "x_python_type_callable.json", + output_path=output_file, + input_file_type=None, + assert_func=assert_file_content, + extra_args=["--output-model-type", "typing.TypedDict"], + ) + + +def test_x_python_type_callable_anyof(output_file: Path) -> None: + """Test x-python-type in anyOf preserves the Callable type.""" + run_main_and_assert( + input_path=JSON_SCHEMA_DATA_PATH / "x_python_type_callable_anyof.json", + output_path=output_file, + input_file_type=None, + assert_func=assert_file_content, + extra_args=["--output-model-type", "typing.TypedDict"], + ) + + +def test_x_python_type_compatible_set(output_file: Path) -> None: + """Test x-python-type with compatible type (Set) uses existing flag logic.""" + run_main_and_assert( + input_path=JSON_SCHEMA_DATA_PATH / "x_python_type_compatible_set.json", + output_path=output_file, + input_file_type=None, + assert_func=assert_file_content, + extra_args=["--output-model-type", "typing.TypedDict"], + ) + + +def test_x_python_type_fqpath(output_file: Path) -> None: + """Test x-python-type with fully qualified path (e.g., collections.abc.Callable).""" + run_main_and_assert( + input_path=JSON_SCHEMA_DATA_PATH / "x_python_type_fqpath.json", + output_path=output_file, + input_file_type=None, + assert_func=assert_file_content, + extra_args=["--output-model-type", "typing.TypedDict"], + ) + + +def test_x_python_type_no_schema_type(output_file: Path) -> None: + """Test x-python-type when schema type is not specified.""" + run_main_and_assert( + input_path=JSON_SCHEMA_DATA_PATH / "x_python_type_no_schema_type.json", + output_path=output_file, + input_file_type=None, + assert_func=assert_file_content, + extra_args=["--output-model-type", "typing.TypedDict"], + ) + + +def test_x_python_type_custom_fqpath(output_file: Path) -> None: + """Test x-python-type with custom fully qualified path not in predefined imports.""" + run_main_and_assert( + input_path=JSON_SCHEMA_DATA_PATH / "x_python_type_custom_fqpath.json", + output_path=output_file, + input_file_type=None, + assert_func=assert_file_content, + extra_args=["--output-model-type", "typing.TypedDict"], + )