Support PEP 604 unions (X | Y) in PlutusData fields

pycardano/plutus.py

@@ -4,6 +4,7 @@
 
 import inspect
 import json
+import types
 import typing
 from dataclasses import dataclass, field, fields
 from enum import Enum
@@ -489,6 +490,11 @@ def id_map(cls, skip_constructor=False):
         return "any"
     if cls == IndefiniteList:
         return "list"
+    if isinstance(cls, types.UnionType):
+        # PEP 604 union (e.g. `bytes | int`) has no __origin__, but it is
+        # equivalent to typing.Union[...] — map it identically so the
+        # constructor ID does not depend on which spelling the author used.
+        return "union<" + ",".join(id_map(a) for a in cls.__args__) + ">"
     if hasattr(cls, "__origin__"):
         origin = getattr(cls, "__origin__")
         if origin == list:
@@ -694,7 +700,7 @@ def _dfs(obj):
                         elif (
                             hasattr(f_info.type, "__origin__")
                             and f_info.type.__origin__ is Union
-                        ):
+                        ) or typing.get_origin(f_info.type) is types.UnionType:
                             t_args = f_info.type.__args__
                             found_match = False
                             for t in t_args:

pycardano/serialization.py

@@ -5,6 +5,7 @@
 import json
 import os
 import re
+import types
 import typing
 from collections import OrderedDict, UserList, defaultdict
 from copy import deepcopy
@@ -854,9 +855,10 @@ def _restore_typed_primitive(
             _restore_typed_primitive(key_t, key): _restore_typed_primitive(val_t, val)
             for key, val in v.items()
         }
-    elif hasattr(t, "__origin__") and (
-        t.__origin__ is Union or t.__origin__ is Optional
-    ):
+    elif (
+        hasattr(t, "__origin__")
+        and (t.__origin__ is Union or t.__origin__ is Optional)
+    ) or typing.get_origin(t) is types.UnionType:
         t_args = t.__args__
         for t in t_args:
             try:

test/pycardano/test_plutus.py

@@ -478,6 +478,52 @@ class B(PlutusData):
     )
 
 
+def test_plutus_data_pep604_union():
+    # A PEP 604 union (`X | Y`) is the modern spelling of typing.Union[X, Y];
+    # the two spellings must behave identically, including the automatic
+    # constructor ID, so a datum's encoding never depends on spelling.
+    @dataclass
+    class Classic(PlutusData):
+        CONSTR_ID = 0
+        a: Union[bytes, int]
+
+    @dataclass
+    class Modern(PlutusData):
+        CONSTR_ID = 0
+        a: bytes | int
+
+    assert id_map(Union[bytes, int]) == id_map(bytes | int) == "union<bytes,int>"
+
+    for v in (b"hello", 42):
+        classic = Classic(a=v)
+        modern = Modern(a=v)
+        assert classic.to_cbor() == modern.to_cbor()
+        assert Modern.from_cbor(modern.to_cbor()) == modern
+        assert Classic.from_cbor(modern.to_cbor()).a == v
+
+
+def test_plutus_data_pep604_union_of_plutus_data():
+    @dataclass
+    class A(PlutusData):
+        CONSTR_ID = 0
+        x: int
+
+    @dataclass
+    class B(PlutusData):
+        CONSTR_ID = 1
+        y: bytes
+
+    @dataclass
+    class Modern(PlutusData):
+        CONSTR_ID = 2
+        u: A | B
+
+    for v in (A(x=1), B(y=b"z")):
+        datum = Modern(u=v)
+        assert Modern.from_cbor(datum.to_cbor()) == datum
+        assert Modern.from_dict(datum.to_dict()) == datum
+
+
 def test_plutus_data_long_bytes():
     @dataclass
     class A(PlutusData):