Fix x-python-type for Optional container types in anyOf schemas (#2840)

Author: koxudaxi

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."""

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):

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:"],
+    )