Refactor default handling for Union and dict types in model generation (#2595)

Author: koxudaxi

src/datamodel_code_generator/model/msgspec.py

@@ -332,8 +332,8 @@ def _get_default_as_struct_model(self) -> str | None:
         """Convert default value to Struct model using msgspec convert."""
         for data_type in self.data_type.data_types or (self.data_type,):
             # TODO: Check nested data_types
-            if data_type.is_dict or self.data_type.is_union:
-                # TODO: Parse Union and dict model for default
+            if data_type.is_dict:
+                # TODO: Parse dict model for default
                 continue  # pragma: no cover
             if data_type.is_list and len(data_type.data_types) == 1:
                 data_type_child = data_type.data_types[0]
@@ -347,6 +347,11 @@ def _get_default_as_struct_model(self) -> str | None:
                         f"type=list[{data_type_child.alias or data_type_child.reference.source.class_name}])"
                     )
             elif data_type.reference and isinstance(data_type.reference.source, Struct):
+                if self.data_type.is_union:
+                    if not isinstance(self.default, (dict, list)):
+                        continue
+                    if isinstance(self.default, dict) and any(dt.is_dict for dt in self.data_type.data_types):
+                        continue
                 return (
                     f"lambda: {self._PARSE_METHOD}({self.default!r},  "
                     f"type={data_type.alias or data_type.reference.source.class_name})"

src/datamodel_code_generator/model/pydantic/base_model.py

@@ -114,8 +114,8 @@ def _get_strict_field_constraint_value(self, constraint: str, value: Any) -> Any
     def _get_default_as_pydantic_model(self) -> str | None:
         for data_type in self.data_type.data_types or (self.data_type,):
             # TODO: Check nested data_types
-            if data_type.is_dict or self.data_type.is_union:
-                # TODO: Parse Union and dict model for default
+            if data_type.is_dict:
+                # TODO: Parse dict model for default
                 continue
             if data_type.is_list and len(data_type.data_types) == 1:
                 data_type_child = data_type.data_types[0]
@@ -128,7 +128,12 @@ def _get_default_as_pydantic_model(self) -> str | None:
                         f"lambda :[{data_type_child.alias or data_type_child.reference.source.class_name}."
                         f"{self._PARSE_METHOD}(v) for v in {self.default!r}]"
                     )
-            elif data_type.reference and isinstance(data_type.reference.source, BaseModelBase):  # pragma: no cover
+            elif data_type.reference and isinstance(data_type.reference.source, BaseModelBase):
+                if self.data_type.is_union:
+                    if not isinstance(self.default, (dict, list)):
+                        continue
+                    if isinstance(self.default, dict) and any(dt.is_dict for dt in self.data_type.data_types):
+                        continue
                 return (
                     f"lambda :{data_type.alias or data_type.reference.source.class_name}."
                     f"{self._PARSE_METHOD}({self.default!r})"

tests/data/expected/main/openapi/msgspec_union_default_object.py

@@ -0,0 +1,22 @@
+# generated by datamodel-codegen:
+#   filename:  union_default_object.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Dict, Union
+
+from msgspec import UNSET, Struct, UnsetType, convert, field
+
+
+class Interval(Struct):
+    start: Union[int, UnsetType] = UNSET
+    end: Union[int, UnsetType] = UNSET
+
+
+class Container(Struct):
+    interval_or_string: Union[Interval, str, UnsetType] = field(
+        default_factory=lambda: convert({'start': 2009, 'end': 2019}, type=Interval)
+    )
+    string_or_interval: Union[Interval, str, UnsetType] = 'some string value'
+    dict_or_interval: Union[Dict[str, str], Interval, UnsetType] = {'key': 'value'}

tests/data/expected/main/openapi/pydantic_v2_union_default_object.py

@@ -0,0 +1,22 @@
+# generated by datamodel-codegen:
+#   filename:  union_default_object.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Dict, Optional, Union
+
+from pydantic import BaseModel, Field
+
+
+class Interval(BaseModel):
+    start: Optional[int] = None
+    end: Optional[int] = None
+
+
+class Container(BaseModel):
+    interval_or_string: Optional[Union[Interval, str]] = Field(
+        default_factory=lambda: Interval.model_validate({'start': 2009, 'end': 2019})
+    )
+    string_or_interval: Optional[Union[Interval, str]] = 'some string value'
+    dict_or_interval: Optional[Union[Dict[str, str], Interval]] = {'key': 'value'}

tests/data/expected/main/openapi/union_default_object.py

@@ -0,0 +1,22 @@
+# generated by datamodel-codegen:
+#   filename:  union_default_object.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Dict, Optional, Union
+
+from pydantic import BaseModel, Field
+
+
+class Interval(BaseModel):
+    start: Optional[int] = None
+    end: Optional[int] = None
+
+
+class Container(BaseModel):
+    interval_or_string: Optional[Union[Interval, str]] = Field(
+        default_factory=lambda: Interval.parse_obj({'start': 2009, 'end': 2019})
+    )
+    string_or_interval: Optional[Union[Interval, str]] = 'some string value'
+    dict_or_interval: Optional[Union[Dict[str, str], Interval]] = {'key': 'value'}

tests/data/openapi/union_default_object.yaml

@@ -0,0 +1,47 @@
+openapi: 3.1.0
+info:
+  title: Union Default Object Test
+  version: 0.1.0
+servers:
+- url: http://example.com
+paths:
+  /test:
+    get:
+      responses:
+        '200':
+          description: OK
+components:
+  schemas:
+    Interval:
+      type: object
+      properties:
+        start:
+          type: integer
+        end:
+          type: integer
+    Container:
+      type: object
+      properties:
+        # Union[Interval, str] with dict default - should use default_factory
+        interval_or_string:
+          anyOf:
+            - $ref: '#/components/schemas/Interval'
+            - type: string
+          default:
+            start: 2009
+            end: 2019
+        # Union[Interval, str] with string default - should NOT use default_factory
+        string_or_interval:
+          anyOf:
+            - $ref: '#/components/schemas/Interval'
+            - type: string
+          default: "some string value"
+        # Union[Dict, Interval] with dict default - should NOT use default_factory (dict arm)
+        dict_or_interval:
+          anyOf:
+            - type: object
+              additionalProperties:
+                type: string
+            - $ref: '#/components/schemas/Interval'
+          default:
+            key: "value"

tests/main/openapi/test_main_openapi.py

@@ -1577,6 +1577,38 @@ def test_main_openapi_default_object(output_model: str, expected_output: str, tm
     )
 
 
+@pytest.mark.parametrize(
+    ("output_model", "expected_output"),
+    [
+        (
+            "pydantic.BaseModel",
+            "union_default_object.py",
+        ),
+        (
+            "pydantic_v2.BaseModel",
+            "pydantic_v2_union_default_object.py",
+        ),
+        (
+            "msgspec.Struct",
+            "msgspec_union_default_object.py",
+        ),
+    ],
+)
+@pytest.mark.skipif(
+    black.__version__.split(".")[0] == "19",
+    reason="Installed black doesn't support the old style",
+)
+def test_main_openapi_union_default_object(output_model: str, expected_output: str, output_file: Path) -> None:
+    """Test OpenAPI generation with Union type default object values."""
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "union_default_object.yaml",
+        output_path=output_file,
+        expected_file=EXPECTED_OPENAPI_PATH / expected_output,
+        input_file_type="openapi",
+        extra_args=["--output-model", output_model, "--target-python-version", "3.9", "--openapi-scopes", "schemas"],
+    )
+
+
 def test_main_dataclass(output_file: Path) -> None:
     """Test OpenAPI generation with dataclass output."""
     run_main_and_assert(