Fix msgspec type hint generation for oneOf/anyOf with null type (#2629)

* Add support for oneOf containing null type in msgspec Struct generation

* Refactor annotated property to improve ClassVar handling and simplify type hint generation

* Add MSGSPEC_LEGACY_BLACK_SKIP decorator to oneOf null type tests

* Replace MSGSPEC_LEGACY_BLACK_SKIP decorator with LEGACY_BLACK_SKIP in oneOf null type tests

Author: koxudaxi

src/datamodel_code_generator/model/msgspec.py

@@ -180,6 +180,19 @@ def get_neither_required_nor_nullable_type(type_: str, use_union_operator: bool)
     return f"{UNION_PREFIX}{type_}{UNION_DELIMITER}{UNSET_TYPE}]"
 
 
+@lru_cache
+def _add_unset_type(type_: str, use_union_operator: bool) -> str:  # noqa: FBT001
+    """Add UnsetType to a type hint without removing None."""
+    if use_union_operator:
+        return f"{type_}{UNION_OPERATOR_DELIMITER}{UNSET_TYPE}"
+    if type_.startswith(UNION_PREFIX):
+        return f"{type_[:-1]}{UNION_DELIMITER}{UNSET_TYPE}]"
+    if type_.startswith(OPTIONAL_PREFIX):  # pragma: no cover
+        inner_type = type_[len(OPTIONAL_PREFIX) : -1]
+        return f"{UNION_PREFIX}{inner_type}{UNION_DELIMITER}{NONE}{UNION_DELIMITER}{UNSET_TYPE}]"
+    return f"{UNION_PREFIX}{type_}{UNION_DELIMITER}{UNSET_TYPE}]"
+
+
 @import_extender
 class DataModelField(DataModelFieldBase):
     """Field implementation for msgspec Struct models."""
@@ -281,7 +294,9 @@ def __str__(self) -> str:
     def type_hint(self) -> str:
         """Return the type hint, using UnsetType for non-required non-nullable fields."""
         type_hint = super().type_hint
-        if self._not_required and not self.nullable and not self.data_type.is_optional:
+        if self._not_required and not self.nullable:
+            if self.data_type.is_optional:
+                return _add_unset_type(type_hint, self.data_type.use_union_operator)
             return get_neither_required_nor_nullable_type(type_hint, self.data_type.use_union_operator)
         return type_hint
 
@@ -316,7 +331,7 @@ def _get_meta_string(self) -> str | None:
         return f"Meta({', '.join(meta_arguments)})" if meta_arguments else None
 
     @property
-    def annotated(self) -> str | None:
+    def annotated(self) -> str | None:  # noqa: PLR0911
         """Get Annotated type hint with Meta constraints.
 
         For ClassVar fields (discriminator tag_field), ClassVar is required
@@ -325,24 +340,26 @@ def annotated(self) -> str | None:
         if self.extras.get("is_classvar"):
             meta = self._get_meta_string()
             if self.use_annotated and meta:
-                annotated_type = f"Annotated[{self.type_hint}, {meta}]"
-                return f"ClassVar[{annotated_type}]"
+                return f"ClassVar[Annotated[{self.type_hint}, {meta}]]"
             return f"ClassVar[{self.type_hint}]"
 
         if not self.use_annotated:  # pragma: no cover
             return None
 
         meta = self._get_meta_string()
-
         if not meta:
             return None
 
-        if not self.required:
-            type_hint = self.data_type.type_hint
-            annotated_type = f"Annotated[{type_hint}, {meta}]"
-            return get_neither_required_nor_nullable_type(annotated_type, self.data_type.use_union_operator)
-
-        return f"Annotated[{self.type_hint}, {meta}]"
+        if self.required:
+            return f"Annotated[{self.type_hint}, {meta}]"
+
+        type_hint = self.data_type.type_hint
+        annotated_type = f"Annotated[{type_hint}, {meta}]"
+        if self.nullable:  # pragma: no cover
+            return annotated_type
+        if self.data_type.is_optional:  # pragma: no cover
+            return _add_unset_type(annotated_type, self.data_type.use_union_operator)
+        return get_neither_required_nor_nullable_type(annotated_type, self.data_type.use_union_operator)
 
     @property
     def needs_annotated_import(self) -> bool:

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

@@ -0,0 +1,29 @@
+# generated by datamodel-codegen:
+#   filename:  msgspec_oneof_with_null.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Annotated, Union
+
+from msgspec import UNSET, Meta, Struct, UnsetType
+from typing_extensions import TypeAlias
+
+OptionalOneofWithNullAndConstraint: TypeAlias = Annotated[str, Meta(max_length=100)]
+
+
+class Model(Struct):
+    required_field: str
+    optional_oneof_with_null: Union[str, None, UnsetType] = UNSET
+    optional_anyof_with_null: Union[str, None, UnsetType] = UNSET
+    optional_field_not_nullable: Union[str, UnsetType] = UNSET
+    optional_oneof_with_null_and_constraint: Union[
+        OptionalOneofWithNullAndConstraint, None, UnsetType
+    ] = UNSET
+    optional_nullable_field: Union[str, UnsetType] = UNSET
+    optional_nullable_with_constraint: Union[
+        Annotated[str, Meta(max_length=50)], UnsetType
+    ] = UNSET
+    optional_nullable_with_min_length: Union[
+        Annotated[str, Meta(min_length=5)], UnsetType
+    ] = UNSET

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

@@ -0,0 +1,29 @@
+# generated by datamodel-codegen:
+#   filename:  msgspec_oneof_with_null.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Annotated
+
+from msgspec import UNSET, Meta, Struct, UnsetType
+from typing_extensions import TypeAlias
+
+OptionalOneofWithNullAndConstraint: TypeAlias = Annotated[str, Meta(max_length=100)]
+
+
+class Model(Struct):
+    required_field: str
+    optional_oneof_with_null: str | None | UnsetType = UNSET
+    optional_anyof_with_null: str | None | UnsetType = UNSET
+    optional_field_not_nullable: str | UnsetType = UNSET
+    optional_oneof_with_null_and_constraint: (
+        OptionalOneofWithNullAndConstraint | None | UnsetType
+    ) = UNSET
+    optional_nullable_field: str | UnsetType = UNSET
+    optional_nullable_with_constraint: (
+        Annotated[str, Meta(max_length=50)] | UnsetType
+    ) = UNSET
+    optional_nullable_with_min_length: (
+        Annotated[str, Meta(min_length=5)] | UnsetType
+    ) = UNSET

tests/data/openapi/msgspec_oneof_with_null.yaml

@@ -0,0 +1,39 @@
+openapi: "3.0.0"
+info:
+  version: 1.0.0
+  title: Test oneOf with null
+components:
+  schemas:
+    Model:
+      type: object
+      properties:
+        required_field:
+          type: string
+        optional_oneof_with_null:
+          oneOf:
+            - type: string
+            - type: "null"
+        optional_anyof_with_null:
+          anyOf:
+            - type: string
+            - type: "null"
+        optional_field_not_nullable:
+          type: string
+        optional_oneof_with_null_and_constraint:
+          oneOf:
+            - type: string
+              maxLength: 100
+            - type: "null"
+        optional_nullable_field:
+          type: string
+          nullable: true
+        optional_nullable_with_constraint:
+          type: string
+          nullable: true
+          maxLength: 50
+        optional_nullable_with_min_length:
+          type: string
+          nullable: true
+          minLength: 5
+      required:
+        - required_field

tests/main/openapi/test_main_openapi.py

@@ -2391,6 +2391,39 @@ def test_main_openapi_msgspec_anyof(min_version: str, output_file: Path) -> None
     )
 
 
+@LEGACY_BLACK_SKIP
+def test_main_openapi_msgspec_oneof_with_null(output_file: Path) -> None:
+    """Test msgspec Struct generation with oneOf containing null type."""
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "msgspec_oneof_with_null.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file="msgspec_oneof_with_null.py",
+        extra_args=[
+            "--output-model-type",
+            "msgspec.Struct",
+        ],
+    )
+
+
+@LEGACY_BLACK_SKIP
+def test_main_openapi_msgspec_oneof_with_null_union_operator(output_file: Path) -> None:
+    """Test msgspec Struct generation with oneOf containing null type using union operator."""
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "msgspec_oneof_with_null.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file="msgspec_oneof_with_null_union_operator.py",
+        extra_args=[
+            "--output-model-type",
+            "msgspec.Struct",
+            "--use-union-operator",
+        ],
+    )
+
+
 def test_main_openapi_referenced_default(output_file: Path) -> None:
     """Test OpenAPI generation with referenced default values."""
     run_main_and_assert(