refactor: simplify get_native_to_arc4_serializer implementation

Author: daniel-makerx

src/_algopy_testing/serialize.py

@@ -1,4 +1,5 @@
 import dataclasses
+import functools
 import inspect
 import typing
 from collections.abc import Callable, Sequence
@@ -15,7 +16,6 @@
 
 @dataclasses.dataclass(frozen=True)
 class _Serializer(typing.Generic[_T, _U]):
-    native_type: type[_T]
     arc4_type: type[_U]
     native_to_arc4: Callable[[_T], _U]
     arc4_to_native: Callable[[_U], _T]
@@ -25,111 +25,92 @@ def identity(i: _T) -> _T:
     return i
 
 
-def get_native_to_arc4_serializer(typ: type[_T]) -> _Serializer:  # type: ignore[type-arg]  # noqa: PLR0911, PLR0912
+def get_native_to_arc4_serializer(typ: type) -> _Serializer[typing.Any, typing.Any]:
     from _algopy_testing import arc4
-    from _algopy_testing.models import Account
-    from _algopy_testing.primitives import BigUInt, Bytes, ImmutableArray, String
+    from _algopy_testing.primitives import ImmutableArray
     from _algopy_testing.protocols import UInt64Backed
 
-    origin_typ = typing.get_origin(typ)
-
-    if inspect.isclass(typ) and issubclass(typ, arc4._ABIEncoded):
-        return _Serializer(
-            native_type=typ, arc4_type=typ, native_to_arc4=identity, arc4_to_native=identity
-        )
-    # For types that are expected to be simple classes for these specific checks
-    if inspect.isclass(typ):
-        if issubclass(typ, bool):
-            return _Serializer(
-                native_type=typ,
-                arc4_type=arc4.Bool,
-                native_to_arc4=arc4.Bool,
-                arc4_to_native=lambda n: n.native,
-            )
+    origin_type = typing.get_origin(typ)
+    if origin_type is tuple:
+        return _get_tuple_serializer(typing.get_args(typ))
+    elif isinstance(typ, type):
+        if issubclass(typ, arc4._ABIEncoded):
+            return _Serializer(arc4_type=typ, native_to_arc4=identity, arc4_to_native=identity)
+        for native_type, simple_arc4_type in _simple_native_to_arc4_type_map().items():
+            if issubclass(typ, native_type):
+                return _Serializer(
+                    arc4_type=simple_arc4_type,
+                    native_to_arc4=simple_arc4_type,
+                    arc4_to_native=lambda n: n.native,
+                )
         if issubclass(typ, UInt64Backed):
             return _Serializer(
-                native_type=typ,
                 arc4_type=arc4.UInt64,
                 native_to_arc4=lambda n: arc4.UInt64(n.int_),
                 arc4_to_native=lambda a: typ.from_int(a.native),
             )
-        if issubclass(typ, BigUInt):
-            return _Serializer(
-                native_type=typ,
-                arc4_type=arc4.UInt512,
-                native_to_arc4=arc4.UInt512,
-                arc4_to_native=lambda a: a.native,
-            )
-        if issubclass(typ, Account):
-            return _Serializer(
-                native_type=typ,
-                arc4_type=arc4.Address,
-                native_to_arc4=arc4.Address,
-                arc4_to_native=lambda a: a.native,
-            )
-        if issubclass(typ, UInt64):
-            return _Serializer(
-                native_type=typ,
-                arc4_type=arc4.UInt64,
-                native_to_arc4=arc4.UInt64,
-                arc4_to_native=lambda a: a.native,
-            )
-        if issubclass(typ, Bytes):
-            return _Serializer(
-                native_type=typ,
-                arc4_type=arc4.DynamicBytes,
-                native_to_arc4=arc4.DynamicBytes,
-                arc4_to_native=lambda a: a.native,
-            )
-        if issubclass(typ, String):
+        if typing.NamedTuple in getattr(typ, "__orig_bases__", []):
+            tuple_fields = tuple(inspect.get_annotations(typ).values())
+            if any(isinstance(f, str) for f in tuple_fields):
+                raise TypeError("string annotations in typing.NamedTuple fields are not supported")
+            return _get_tuple_serializer(tuple_fields)
+        if issubclass(typ, ImmutableArray):
+            native_element_type = typ._element_type
+            element_serializer = get_native_to_arc4_serializer(native_element_type)
+            arc4_element_type = element_serializer.arc4_type
+            arc4_type = arc4.DynamicArray[arc4_element_type]  # type: ignore[valid-type]
             return _Serializer(
-                native_type=typ,
-                arc4_type=arc4.String,
-                native_to_arc4=arc4.String,
-                arc4_to_native=lambda a: a.native,
+                arc4_type=arc4_type,
+                native_to_arc4=lambda arr: arc4_type(
+                    *(element_serializer.native_to_arc4(e) for e in arr)
+                ),
+                arc4_to_native=lambda arr: typ(
+                    *(element_serializer.arc4_to_native(e) for e in arr)
+                ),
             )
-
-    if origin_typ is tuple or (inspect.isclass(typ) and issubclass(typ, tuple)):
-        if typing.NamedTuple in getattr(typ, "__orig_bases__", []):
-            tuple_fields: Sequence[type] = list(inspect.get_annotations(typ).values())
-        else:
-            tuple_fields = typing.get_args(typ)
-        serializers = [get_native_to_arc4_serializer(i) for i in tuple_fields]
-
-        def _items_to_arc4(items: Sequence[object]) -> tuple[object, ...]:
-            result = []
-            for item, serializer in zip(items, serializers, strict=True):
-                result.append(serializer.native_to_arc4(item))
-            return tuple(result)
-
-        def _items_to_native(items: Sequence[object]) -> tuple[object, ...]:
-            result = []
-            for item, serializer in zip(items, serializers, strict=True):
-                result.append(serializer.arc4_to_native(item))
-            return tuple(result)
-
-        return _Serializer(
-            native_type=typ,
-            arc4_type=arc4.Tuple[*(s.arc4_type for s in serializers)],  # type: ignore[misc]
-            native_to_arc4=lambda t: arc4.Tuple(_items_to_arc4(t)),
-            arc4_to_native=lambda t: _items_to_native(t),
-        )
-    if issubclass(typ, ImmutableArray):
-        native_element_type = typ._element_type
-        element_serializer = get_native_to_arc4_serializer(native_element_type)
-        arc4_element_type = element_serializer.arc4_type
-        arc4_type = arc4.DynamicArray[arc4_element_type]  # type: ignore[valid-type]
-        return _Serializer(
-            native_type=typ,
-            arc4_type=arc4_type,
-            native_to_arc4=lambda arr: arc4_type(
-                *(element_serializer.native_to_arc4(e) for e in arr)
-            ),
-            arc4_to_native=lambda arr: typ(*(element_serializer.arc4_to_native(e) for e in arr)),
-        )
     raise TypeError(f"unserializable type: {typ}")
 
 
+@functools.cache
+def _simple_native_to_arc4_type_map() -> dict[type, type]:
+    from _algopy_testing import arc4
+    from _algopy_testing.models import Account
+    from _algopy_testing.primitives import BigUInt, Bytes, String
+
+    return {
+        bool: arc4.Bool,
+        UInt64: arc4.UInt64,
+        BigUInt: arc4.UInt512,
+        Account: arc4.Address,
+        Bytes: arc4.DynamicBytes,
+        String: arc4.String,
+    }
+
+
+def _get_tuple_serializer(item_types: tuple[type, ...]) -> _Serializer[typing.Any, typing.Any]:
+    from _algopy_testing import arc4
+
+    serializers = [get_native_to_arc4_serializer(i) for i in item_types]
+
+    def _items_to_arc4(items: Sequence[object]) -> tuple[object, ...]:
+        result = []
+        for item, serializer in zip(items, serializers, strict=True):
+            result.append(serializer.native_to_arc4(item))
+        return tuple(result)
+
+    def _items_to_native(items: Sequence[object]) -> tuple[object, ...]:
+        result = []
+        for item, serializer in zip(items, serializers, strict=True):
+            result.append(serializer.arc4_to_native(item))
+        return tuple(result)
+
+    return _Serializer(
+        arc4_type=arc4.Tuple[*(s.arc4_type for s in serializers)],  # type: ignore[misc]
+        native_to_arc4=lambda t: arc4.Tuple(_items_to_arc4(t)),
+        arc4_to_native=lambda t: _items_to_native(t),
+    )
+
+
 def serialize_to_bytes(value: object) -> bytes:
     return native_to_arc4(value).bytes.value