Merge pull request #36 from algorandfoundation/feat/size-of

feat: add stub implementation of `algopy.size_of` function

Author: boblat

docs/coverage.md

@@ -3,7 +3,7 @@
 See which `algorand-python` stubs are implemented by the `algorand-python-testing` library. See the [Concepts](testing-guide/concepts.md#types-of-algopy-stub-implementations) section for more details on the implementation categories. Refer to the [`algorand-python` stubs API](api.md) for the full list of the stubs for which the `algorand-python-testing` library provides implementations referenced in the table below.
 
 | Name                                        | Implementation type |
-| ------------------------------------------- | ------------------- |
+|---------------------------------------------|---------------------|
 | algopy.Account                              | Emulated            |
 | algopy.Application                          | Emulated            |
 | algopy.Asset                                | Emulated            |
@@ -33,6 +33,7 @@ See which `algorand-python` stubs are implemented by the `algorand-python-testin
 | algopy.ensure_budget                        | Emulated            |
 | algopy.log                                  | Emulated            |
 | algopy.logicsig                             | Emulated            |
+| algopy.size_of                              | Emulated            |
 | algopy.subroutine                           | Native              |
 | algopy.uenumerate                           | Native              |
 | algopy.urange                               | Native              |

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,22 @@
+import types
+
+from _algopy_testing.primitives.uint64 import UInt64
+
+
+def size_of(typ: type | object, /) -> UInt64:
+    from _algopy_testing.arc4 import get_max_bytes_static_len
+    from _algopy_testing.serialize import get_native_to_arc4_serializer
+
+    if isinstance(typ, types.GenericAlias):
+        pass
+    elif not isinstance(typ, type):
+        typ = type(typ)
+
+    if typ is bool:  # treat bool on its own as a uint64
+        typ = UInt64
+    serializer = get_native_to_arc4_serializer(typ)  # type: ignore[arg-type]
+    type_info = serializer.arc4_type._type_info
+    size = get_max_bytes_static_len(type_info)
+    if size is None:
+        raise ValueError(f"{typ} is dynamically sized")
+    return UInt64(size)

src/algopy/__init__.py

@@ -17,7 +17,7 @@
 from _algopy_testing.primitives import Array, BigUInt, Bytes, ImmutableArray, String, UInt64
 from _algopy_testing.protocols import BytesBacked
 from _algopy_testing.state import Box, BoxMap, BoxRef, GlobalState, LocalState
-from _algopy_testing.utilities import OpUpFeeSource, ensure_budget, log
+from _algopy_testing.utilities import OpUpFeeSource, ensure_budget, log, size_of
 
 from . import arc4, gtxn, itxn, op
 
@@ -59,6 +59,7 @@
     "logicsig",
     "op",
     "subroutine",
+    "size_of",
     "uenumerate",
     "urange",
 ]

tests/utilities/test_size_of.py

@@ -0,0 +1,65 @@
+import typing
+
+import pytest
+from algopy import Account, Application, Asset, Bytes, String, UInt64, arc4, size_of
+
+
+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
+
+
+@pytest.mark.parametrize(
+    "typ",
+    [
+        arc4.StaticArray[arc4.DynamicBytes, typing.Literal[7]],
+        tuple[arc4.DynamicBytes, Bytes],
+        arc4.Tuple[arc4.UInt64, arc4.String],
+        MyDynamicSizedTuple,
+    ],
+)
+def test_size_of_dynamic(typ: type) -> None:
+    with pytest.raises(ValueError, match="is dynamically sized"):
+        size_of(typ)