Support PEP 604 (X | Y) union notation (lovasoa#219)

Fixes lovasoa#194

Note that the PEP 604 notation is only supported in Python version
3.10 and above.

Author: dairiki

marshmallow_dataclass/__init__.py

@@ -495,8 +495,8 @@ def _field_for_generic_type(
     If the type is a generic interface, resolve the arguments and construct the appropriate Field.
     """
     origin = typing_inspect.get_origin(typ)
+    arguments = typing_inspect.get_args(typ, True)
     if origin:
-        arguments = typing_inspect.get_args(typ, True)
         # Override base_schema.TYPE_MAPPING to change the class used for generic types below
         type_mapping = base_schema.TYPE_MAPPING if base_schema else {}
 
@@ -557,38 +557,38 @@ def _field_for_generic_type(
                 ),
                 **metadata,
             )
-        elif typing_inspect.is_union_type(typ):
-            if typing_inspect.is_optional_type(typ):
-                metadata["allow_none"] = metadata.get("allow_none", True)
-                metadata["dump_default"] = metadata.get("dump_default", None)
-                if not metadata.get("required"):
-                    metadata["load_default"] = metadata.get("load_default", None)
-                metadata.setdefault("required", False)
-            subtypes = [t for t in arguments if t is not NoneType]  # type: ignore
-            if len(subtypes) == 1:
-                return field_for_schema(
-                    subtypes[0],
-                    metadata=metadata,
-                    base_schema=base_schema,
-                    typ_frame=typ_frame,
-                )
-            from . import union_field
+    if typing_inspect.is_union_type(typ):
+        if typing_inspect.is_optional_type(typ):
+            metadata["allow_none"] = metadata.get("allow_none", True)
+            metadata["dump_default"] = metadata.get("dump_default", None)
+            if not metadata.get("required"):
+                metadata["load_default"] = metadata.get("load_default", None)
+            metadata.setdefault("required", False)
+        subtypes = [t for t in arguments if t is not NoneType]  # type: ignore
+        if len(subtypes) == 1:
+            return field_for_schema(
+                subtypes[0],
+                metadata=metadata,
+                base_schema=base_schema,
+                typ_frame=typ_frame,
+            )
+        from . import union_field
 
-            return union_field.Union(
-                [
-                    (
+        return union_field.Union(
+            [
+                (
+                    subtyp,
+                    field_for_schema(
                         subtyp,
-                        field_for_schema(
-                            subtyp,
-                            metadata={"required": True},
-                            base_schema=base_schema,
-                            typ_frame=typ_frame,
-                        ),
-                    )
-                    for subtyp in subtypes
-                ],
-                **metadata,
-            )
+                        metadata={"required": True},
+                        base_schema=base_schema,
+                        typ_frame=typ_frame,
+                    ),
+                )
+                for subtyp in subtypes
+            ],
+            **metadata,
+        )
     return None
 
 

tests/test_union.py

@@ -1,4 +1,5 @@
 from dataclasses import field
+import sys
 import unittest
 from typing import List, Optional, Union, Dict
 
@@ -180,3 +181,18 @@ class IntOrStrWithDefault:
             )
         with self.assertRaises(marshmallow.exceptions.ValidationError):
             schema.load({"value": None})
+
+    @unittest.skipIf(sys.version_info < (3, 10), "No PEP604 support in py<310")
+    def test_pep604_union(self):
+        @dataclass
+        class PEP604IntOrStr:
+            value: int | str
+
+        schema = PEP604IntOrStr.Schema()
+        data_in = {"value": "hello"}
+        loaded = schema.load(data_in)
+        self.assertEqual(loaded, PEP604IntOrStr(value="hello"))
+        self.assertEqual(schema.dump(loaded), data_in)
+
+        data_in = {"value": 42}
+        self.assertEqual(schema.dump(schema.load(data_in)), data_in)