diff --git a/pycardano/plutus.py b/pycardano/plutus.py index b6b2ac3a..7e71c508 100644 --- a/pycardano/plutus.py +++ b/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: diff --git a/pycardano/serialization.py b/pycardano/serialization.py index 2422b3f3..3e016e8d 100644 --- a/pycardano/serialization.py +++ b/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: diff --git a/test/pycardano/test_plutus.py b/test/pycardano/test_plutus.py index 49eb1038..151f005e 100644 --- a/test/pycardano/test_plutus.py +++ b/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" + + 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):