diff --git a/src/datamodel_code_generator/parser/jsonschema.py b/src/datamodel_code_generator/parser/jsonschema.py index 311c5fa17..37f7828e0 100644 --- a/src/datamodel_code_generator/parser/jsonschema.py +++ b/src/datamodel_code_generator/parser/jsonschema.py @@ -1236,8 +1236,6 @@ class decorator which does not preserve staticmethod descriptors. if not x_python_type or not isinstance(x_python_type, str): return {} - base_type = x_python_type.split("[")[0].strip() - type_to_flag: dict[str, dict[str, bool]] = { "Set": {"is_set": True}, "FrozenSet": {"is_frozen_set": True}, @@ -1249,7 +1247,34 @@ class decorator which does not preserve staticmethod descriptors. "MutableSet": {"is_set": True}, } - return type_to_flag.get(base_type, {}) + base_type = x_python_type.split("[")[0].strip() + if base_type in type_to_flag: + return type_to_flag[base_type] + + if base_type in {"Union", "Optional"}: + bracket_start = x_python_type.find("[") + if bracket_start != -1: + inner = x_python_type[bracket_start + 1 : -1] + depth = 0 + current = "" + for char in inner: + if char == "[": + depth += 1 + elif char == "]": + depth -= 1 + if char == "," and depth == 0: + arg_base = current.strip().split("[")[0] + if arg_base in type_to_flag: + return type_to_flag[arg_base] + current = "" + else: + current += char + if current.strip(): + arg_base = current.strip().split("[")[0] + if arg_base in type_to_flag: + return type_to_flag[arg_base] + + return {} def _apply_title_as_name(self, name: str, obj: JsonSchemaObject) -> str: """Apply title as name if use_title_as_name is enabled.""" diff --git a/tests/data/python/input_model/pydantic_models.py b/tests/data/python/input_model/pydantic_models.py index 79725f007..5bc4fed1a 100644 --- a/tests/data/python/input_model/pydantic_models.py +++ b/tests/data/python/input_model/pydantic_models.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Mapping, Sequence -from typing import FrozenSet, Optional, Set +from typing import FrozenSet, Optional, Set, Union from pydantic import BaseModel @@ -32,6 +32,7 @@ class ModelWithPythonTypes(BaseModel): tag_obj: Tag nested_in_list: list[Set[int]] optional_set: Optional[Set[str]] + nullable_frozenset: Union[None, FrozenSet[str]] class RecursiveNode(BaseModel): diff --git a/tests/test_input_model.py b/tests/test_input_model.py index 7d4b4f705..c99f0dab5 100644 --- a/tests/test_input_model.py +++ b/tests/test_input_model.py @@ -559,3 +559,34 @@ def test_input_model_recursive_model_types(tmp_path: Path) -> None: output_path=tmp_path / "output.py", expected_output_contains=["set[str]", "value:"], ) + + +@SKIP_PYDANTIC_V1 +def test_input_model_optional_set_type(tmp_path: Path) -> None: + """Test that Optional[Set[str]] is preserved when converting Pydantic model.""" + run_input_model_and_assert( + input_model="tests.data.python.input_model.pydantic_models:ModelWithPythonTypes", + output_path=tmp_path / "output.py", + expected_output_contains=["set[str] | None", "optional_set:"], + ) + + +@SKIP_PYDANTIC_V1 +def test_input_model_optional_set_to_typeddict(tmp_path: Path) -> None: + """Test that Optional[Set[str]] works when outputting to TypedDict.""" + run_input_model_and_assert( + input_model="tests.data.python.input_model.pydantic_models:ModelWithPythonTypes", + output_path=tmp_path / "output.py", + extra_args=["--output-model-type", "typing.TypedDict"], + expected_output_contains=["TypedDict", "set[str] | None", "optional_set:"], + ) + + +@SKIP_PYDANTIC_V1 +def test_input_model_union_none_frozenset(tmp_path: Path) -> None: + """Test that Union[None, FrozenSet[str]] is preserved (container not first arg).""" + run_input_model_and_assert( + input_model="tests.data.python.input_model.pydantic_models:ModelWithPythonTypes", + output_path=tmp_path / "output.py", + expected_output_contains=["frozenset[str] | None", "nullable_frozenset:"], + )