feat: add stub implementation of algopy.size_of function for calculating static storage size of types

Author: boblat

src/_algopy_testing/arc4.py

@@ -1256,10 +1256,20 @@ def _compress_multiple_bool(value_list: list[Bool]) -> int:
     return result
 
 
+def get_max_bytes_static_len(type_info: _TypeInfo) -> int | None:
+    return _get_max_bytes_len_impl(type_info, static_size_only=True)
+
+
 def _get_max_bytes_len(type_info: _TypeInfo) -> int:
+    return _get_max_bytes_len_impl(type_info) or 0
+
+
+def _get_max_bytes_len_impl(type_info: _TypeInfo, *, static_size_only: bool = False) -> int | None:
     size = 0
     if isinstance(type_info, _DynamicArrayTypeInfo):
         size += _ABI_LENGTH_SIZE
+        if static_size_only:
+            return None
     elif isinstance(type_info, _TupleTypeInfo | _StructTypeInfo | _StaticArrayTypeInfo):
         i = 0
         if isinstance(type_info, _TupleTypeInfo | _StructTypeInfo):
@@ -1276,12 +1286,18 @@ def _get_max_bytes_len(type_info: _TypeInfo) -> int:
                 if bool_num % 8 != 0:
                     size += 1
             else:
-                child_byte_size = _get_max_bytes_len(child_types[i])
+                child_byte_size = _get_max_bytes_len_impl(
+                    child_types[i], static_size_only=static_size_only
+                )
+                if child_byte_size is None:
+                    return None
                 size += child_byte_size
             i += 1
     elif isinstance(type_info, _UIntTypeInfo):
         size = type_info.max_bytes_len
 
+    elif static_size_only:
+        return None
     return size
 
 

src/_algopy_testing/utilities/__init__.py

@@ -1,4 +1,5 @@
 from _algopy_testing.utilities.budget import OpUpFeeSource, ensure_budget
 from _algopy_testing.utilities.log import log
+from _algopy_testing.utilities.size_of import size_of
 
-__all__ = ["OpUpFeeSource", "ensure_budget", "log"]
+__all__ = ["OpUpFeeSource", "ensure_budget", "log", "size_of"]

src/_algopy_testing/utilities/size_of.py

@@ -0,0 +1,45 @@
+import inspect
+import typing
+
+from _algopy_testing.primitives.uint64 import UInt64
+
+
+def size_of(typ: type | object, /) -> UInt64:
+    from _algopy_testing.arc4 import _TupleTypeInfo, _TypeInfo, get_max_bytes_static_len
+    from _algopy_testing.models.account import Account
+    from _algopy_testing.models.application import Application
+    from _algopy_testing.models.asset import Asset
+    from _algopy_testing.serialize import get_native_to_arc4_serializer
+
+    # Check for types with _type_info attribute
+    if hasattr(typ, "_type_info") or isinstance(typ, _TypeInfo):
+        size = get_max_bytes_static_len(typ if isinstance(typ, _TypeInfo) else typ._type_info)
+        if size is not None:
+            return UInt64(size)
+        else:
+            raise ValueError(f"{typ} is dynamically sized")
+
+    # Fixed-size types
+    fixed_sizes = {UInt64: 8, Account: 32, Application: 8, Asset: 8, bool: 8}
+
+    # Check if type matches any fixed-size type
+    for cls, size in fixed_sizes.items():
+        if isinstance(typ, cls) or typ == cls:
+            return UInt64(size)
+
+    # Handle tuple types
+    is_tuple = (isinstance(typ, type) and issubclass(typ, tuple)) or typing.get_origin(
+        typ
+    ) is tuple
+    if is_tuple:
+        if typing.NamedTuple in getattr(typ, "__orig_bases__", []):
+            tuple_fields = list(
+                inspect.get_annotations(typ if isinstance(typ, type) else typ.__class__).values()
+            )
+        else:
+            tuple_fields = list(typing.get_args(typ))
+
+        type_infos = [get_native_to_arc4_serializer(f).arc4_type._type_info for f in tuple_fields]
+        return size_of(_TupleTypeInfo(type_infos))
+
+    raise ValueError(f"{typ} is dynamically sized")

tests/utilities/test_size_of.py

@@ -0,0 +1,65 @@
+import re
+import typing
+
+import pytest
+from _algopy_testing.utilities.size_of import size_of
+from algopy import Account, Application, Asset, Bytes, String, UInt64, arc4
+
+
+class Swapped(arc4.Struct):
+    a: arc4.UInt64
+    b: arc4.Bool
+    c: arc4.Tuple[arc4.UInt64, arc4.Bool, arc4.Bool]
+    d: arc4.StaticArray[arc4.Bool, typing.Literal[10]]
+    e: arc4.Tuple[arc4.UInt64, arc4.StaticArray[arc4.UInt64, typing.Literal[3]]]
+
+
+class WhatsMySize(typing.NamedTuple):
+    foo: UInt64
+    bar: bool
+    baz: Swapped
+
+
+class MyTuple(typing.NamedTuple):
+    foo: UInt64
+    bar: bool
+    baz: bool
+
+
+class MyDynamicSizedTuple(typing.NamedTuple):
+    foo: UInt64
+    bar: String
+
+
+def test_size_of() -> None:
+    x = arc4.UInt64(0)
+    assert size_of(x) == 8
+    assert size_of(arc4.UInt64) == 8
+    assert size_of(UInt64) == 8
+    assert size_of(arc4.Address) == 32
+    assert size_of(Account) == 32
+    assert size_of(Application) == 8
+    assert size_of(Asset) == 8
+    assert size_of(bool) == 8
+    assert size_of(tuple[bool]) == 1
+    assert size_of(tuple[bool, bool, bool, bool, bool, bool, bool, bool]) == 1
+    assert size_of(tuple[bool, bool, bool, bool, bool, bool, bool, bool, bool]) == 2
+    assert size_of(tuple[arc4.UInt64, UInt64, bool, arc4.Bool]) == 17
+    assert size_of(arc4.Tuple[arc4.UInt64, arc4.Bool, arc4.Bool] == 9)
+    assert size_of(MyTuple) == 9
+    assert size_of(WhatsMySize) == 61
+    assert size_of(arc4.StaticArray[arc4.Byte, typing.Literal[7]]) == 7
+    assert size_of(arc4.StaticArray(arc4.Byte(), arc4.Byte())) == 2
+    assert size_of(Swapped) == 52
+
+    with pytest.raises(ValueError, match=re.compile("is dynamically sized")):
+        size_of(arc4.StaticArray[arc4.DynamicBytes, typing.Literal[7]])
+
+    with pytest.raises(ValueError, match=re.compile("is dynamically sized")):
+        size_of(tuple[arc4.DynamicBytes, Bytes])
+
+    with pytest.raises(ValueError, match=re.compile("is dynamically sized")):
+        size_of(arc4.Tuple[arc4.UInt64, arc4.String])
+
+    with pytest.raises(ValueError, match=re.compile("is dynamically sized")):
+        size_of(MyDynamicSizedTuple)