Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 28 additions & 3 deletions src/datamodel_code_generator/parser/jsonschema.py
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -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."""
Expand Down
3 changes: 2 additions & 1 deletion tests/data/python/input_model/pydantic_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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):
Expand Down
31 changes: 31 additions & 0 deletions tests/test_input_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:"],
)
Loading