Modify ABI DynamicBytes and StaticBytes to support router/subroutine (#…

…514)

CHANGELOG.md

@@ -7,6 +7,7 @@
 
 ## Fixed
 * Fix AST duplication bug in `String.set` when called with an `Expr` argument ([#508](https://github.com/algorand/pyteal/pull/508))
+* Modify ABI `DynamicBytes` and `StaticBytes` to support `router`/`subroutine` ([#514](https://github.com/algorand/pyteal/pull/514))
 
 # 0.16.0
 

pyteal/ast/abi/__init__.py

@@ -35,10 +35,16 @@
     Field,
 )
 from pyteal.ast.abi.array_base import ArrayTypeSpec, Array, ArrayElement
-from pyteal.ast.abi.array_static import StaticArrayTypeSpec, StaticArray, StaticBytes
+from pyteal.ast.abi.array_static import (
+    StaticArrayTypeSpec,
+    StaticArray,
+    StaticBytesTypeSpec,
+    StaticBytes,
+)
 from pyteal.ast.abi.array_dynamic import (
     DynamicArrayTypeSpec,
     DynamicArray,
+    DynamicBytesTypeSpec,
     DynamicBytes,
 )
 from pyteal.ast.abi.reference_type import (
@@ -132,9 +138,11 @@
     "ArrayElement",
     "StaticArrayTypeSpec",
     "StaticArray",
+    "StaticBytesTypeSpec",
     "StaticBytes",
     "DynamicArrayTypeSpec",
     "DynamicArray",
+    "DynamicBytesTypeSpec",
     "DynamicBytes",
     "MethodReturn",
     "Transaction",

pyteal/ast/abi/array_dynamic.py

@@ -18,7 +18,7 @@ class DynamicArrayTypeSpec(ArrayTypeSpec[T]):
     def new_instance(self) -> "DynamicArray[T]":
         return DynamicArray(self)
 
-    def annotation_type(self) -> "type[DynamicArray[T]]":
+    def annotation_type(self) -> type["DynamicArray[T]"]:
         return DynamicArray[self.value_type_spec().annotation_type()]  # type: ignore[misc]
 
     def is_length_dynamic(self) -> bool:
@@ -100,11 +100,34 @@ def length(self) -> Expr:
 DynamicArray.__module__ = "pyteal.abi"
 
 
+class DynamicBytesTypeSpec(DynamicArrayTypeSpec[Byte]):
+    def __init__(self) -> None:
+        super().__init__(ByteTypeSpec())
+
+    def new_instance(self) -> "DynamicBytes":
+        return DynamicBytes()
+
+    def annotation_type(self) -> type["DynamicBytes"]:
+        return DynamicBytes
+
+    def __str__(self) -> str:
+        return "byte[]"
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, DynamicBytesTypeSpec)
+
+
+DynamicBytesTypeSpec.__module__ = "pyteal.abi"
+
+
 class DynamicBytes(DynamicArray[Byte]):
     """The convenience class that represents ABI dynamic byte array."""
 
     def __init__(self) -> None:
-        super().__init__(DynamicArrayTypeSpec(ByteTypeSpec()))
+        super().__init__(DynamicBytesTypeSpec())
+
+    def type_spec(self) -> DynamicBytesTypeSpec:
+        return DynamicBytesTypeSpec()
 
     def set(
         self,

pyteal/ast/abi/array_dynamic_test.py

@@ -20,6 +20,11 @@ def test_DynamicArrayTypeSpec_init():
         assert dynamicArrayType.is_length_dynamic()
         assert dynamicArrayType._stride() == elementType.byte_length_static()
 
+    dynamicBytesType = abi.DynamicBytesTypeSpec()
+    assert isinstance(dynamicBytesType.value_type_spec(), abi.ByteTypeSpec)
+    assert dynamicBytesType.is_length_dynamic()
+    assert dynamicBytesType._stride() == 1
+
     for elementType in DYNAMIC_TYPES:
         dynamicArrayType = abi.DynamicArrayTypeSpec(elementType)
         assert dynamicArrayType.value_type_spec() is elementType
@@ -32,6 +37,8 @@ def test_DynamicArrayTypeSpec_str():
         dynamicArrayType = abi.DynamicArrayTypeSpec(elementType)
         assert str(dynamicArrayType) == "{}[]".format(elementType)
 
+    assert str(abi.DynamicBytesTypeSpec()) == "byte[]"
+
 
 def test_DynamicArrayTypeSpec_new_instance():
     for elementType in STATIC_TYPES + DYNAMIC_TYPES:
@@ -40,26 +47,40 @@ def test_DynamicArrayTypeSpec_new_instance():
         assert isinstance(instance, abi.DynamicArray)
         assert instance.type_spec() == dynamicArrayType
 
+    dynamicBytesType = abi.DynamicBytesTypeSpec()
+    instance = dynamicBytesType.new_instance()
+    assert isinstance(instance, abi.DynamicBytes)
+    assert instance.type_spec() == dynamicBytesType
+
 
 def test_DynamicArrayTypeSpec_eq():
     for elementType in STATIC_TYPES + DYNAMIC_TYPES:
         dynamicArrayType = abi.DynamicArrayTypeSpec(elementType)
         assert dynamicArrayType == dynamicArrayType
         assert dynamicArrayType != abi.TupleTypeSpec(dynamicArrayType)
 
+    dynamicBytesType = abi.DynamicBytesTypeSpec()
+    assert dynamicBytesType == dynamicBytesType
+    assert dynamicBytesType != abi.TupleTypeSpec(dynamicBytesType)
+
 
 def test_DynamicArrayTypeSpec_is_dynamic():
     for elementType in STATIC_TYPES + DYNAMIC_TYPES:
         dynamicArrayType = abi.DynamicArrayTypeSpec(elementType)
         assert dynamicArrayType.is_dynamic()
 
+    assert abi.DynamicBytesTypeSpec().is_dynamic()
+
 
 def test_DynamicArrayTypeSpec_byte_length_static():
     for elementType in STATIC_TYPES + DYNAMIC_TYPES:
         dynamicArrayType = abi.DynamicArrayTypeSpec(elementType)
         with pytest.raises(ValueError):
             dynamicArrayType.byte_length_static()
 
+    with pytest.raises(ValueError):
+        abi.DynamicBytesTypeSpec().byte_length_static()
+
 
 def test_DynamicArray_decode():
     encoded = pt.Bytes("encoded")

pyteal/ast/abi/array_static.py

@@ -13,7 +13,6 @@
 from pyteal.ast.abi.uint import Byte, ByteTypeSpec
 from pyteal.ast.abi.array_base import ArrayTypeSpec, Array, ArrayElement
 
-
 T = TypeVar("T", bound=BaseType)
 N = TypeVar("N", bound=int)
 
@@ -30,7 +29,7 @@ def __init__(self, value_type_spec: TypeSpec, array_length: int) -> None:
     def new_instance(self) -> "StaticArray[T, N]":
         return StaticArray(self)
 
-    def annotation_type(self) -> "type[StaticArray[T, N]]":
+    def annotation_type(self) -> type["StaticArray[T, N]"]:
         return StaticArray[  # type: ignore[misc]
             self.value_spec.annotation_type(), Literal[self.array_length]  # type: ignore
         ]
@@ -149,11 +148,30 @@ def __getitem__(self, index: Union[int, Expr]) -> "ArrayElement[T]":
 StaticArray.__module__ = "pyteal.abi"
 
 
+class StaticBytesTypeSpec(StaticArrayTypeSpec[Byte, N], Generic[N]):
+    def __init__(self, array_length: int) -> None:
+        super().__init__(ByteTypeSpec(), array_length)
+
+    def new_instance(self) -> "StaticBytes[N]":
+        return StaticBytes(self)
+
+    def annotation_type(self) -> type["StaticBytes[N]"]:
+        return StaticBytes[  # type: ignore[misc]
+            Literal[self.array_length]  # type: ignore
+        ]
+
+
+StaticBytesTypeSpec.__module__ = "pyteal.abi"
+
+
 class StaticBytes(StaticArray[Byte, N], Generic[N]):
     """The convenience class that represents ABI static byte array."""
 
-    def __init__(self, static_len: N) -> None:
-        super().__init__(StaticArrayTypeSpec(ByteTypeSpec(), static_len))
+    def __init__(self, array_type_spec: StaticBytesTypeSpec[N]) -> None:
+        super().__init__(array_type_spec)
+
+    def type_spec(self) -> StaticBytesTypeSpec:
+        return cast(StaticBytesTypeSpec[N], super().type_spec())
 
     def set(
         self,

pyteal/ast/abi/array_static_test.py

@@ -23,6 +23,13 @@ def test_StaticArrayTypeSpec_init():
         with pytest.raises(TypeError):
             abi.StaticArrayTypeSpec(elementType, -1)
 
+    for length in range(256):
+        staticBytesType = abi.StaticBytesTypeSpec(length)
+        assert isinstance(staticBytesType.value_type_spec(), abi.ByteTypeSpec)
+        assert not staticBytesType.is_length_dynamic()
+        assert staticBytesType._stride() == 1
+        assert staticBytesType.length_static() == length
+
     for elementType in DYNAMIC_TYPES:
         for length in range(256):
             staticArrayType = abi.StaticArrayTypeSpec(elementType, length)
@@ -41,6 +48,9 @@ def test_StaticArrayTypeSpec_str():
             staticArrayType = abi.StaticArrayTypeSpec(elementType, length)
             assert str(staticArrayType) == "{}[{}]".format(elementType, length)
 
+    for length in range(256):
+        assert str(abi.StaticBytesTypeSpec(length)) == f"byte[{length}]"
+
 
 def test_StaticArrayTypeSpec_new_instance():
     for elementType in STATIC_TYPES + DYNAMIC_TYPES:
@@ -53,6 +63,12 @@ def test_StaticArrayTypeSpec_new_instance():
             )
             assert instance.type_spec() == staticArrayType
 
+    for length in range(256):
+        staticBytesType = abi.StaticBytesTypeSpec(length)
+        instance = staticBytesType.new_instance()
+        assert isinstance(instance, abi.StaticBytes)
+        assert instance.type_spec() == staticBytesType
+
 
 def test_StaticArrayTypeSpec_eq():
     for elementType in STATIC_TYPES + DYNAMIC_TYPES:
@@ -64,13 +80,24 @@ def test_StaticArrayTypeSpec_eq():
                 abi.TupleTypeSpec(elementType), length
             )
 
+    for length in range(256):
+        staticBytesType = abi.StaticBytesTypeSpec(length)
+        assert staticBytesType == staticBytesType
+        assert staticBytesType != abi.StaticBytesTypeSpec(length + 1)
+        assert staticBytesType != abi.StaticArrayTypeSpec(
+            abi.TupleTypeSpec(abi.ByteTypeSpec()), length
+        )
+
 
 def test_StaticArrayTypeSpec_is_dynamic():
     for elementType in STATIC_TYPES:
         for length in range(256):
             staticArrayType = abi.StaticArrayTypeSpec(elementType, length)
             assert not staticArrayType.is_dynamic()
 
+    for length in range(256):
+        assert not abi.StaticBytesTypeSpec(length).is_dynamic()
+
     for elementType in DYNAMIC_TYPES:
         for length in range(256):
             staticArrayType = abi.StaticArrayTypeSpec(elementType, length)
@@ -92,6 +119,13 @@ def test_StaticArrayTypeSpec_byte_length_static():
                 actual == expected
             ), "failed with element type {} and length {}".format(elementType, length)
 
+    for length in range(256):
+        staticBytesType = abi.StaticBytesTypeSpec(length)
+        actual = staticBytesType.byte_length_static()
+        assert (
+            actual == length
+        ), f"failed with element type {staticBytesType.value_type_spec()} and length {length}"
+
     for elementType in DYNAMIC_TYPES:
         for length in range(256):
             staticArrayType = abi.StaticArrayTypeSpec(elementType, length)
@@ -246,7 +280,6 @@ def test_StaticArray_set_computed():
 # AACS key recovery
 BYTE_HEX_TEST_CASE = "09f911029d74e35bd84156c5635688c0"
 
-
 BYTES_SET_TESTCASES = [
     bytes.fromhex(BYTE_HEX_TEST_CASE),
     bytearray.fromhex(BYTE_HEX_TEST_CASE),
@@ -255,7 +288,7 @@ def test_StaticArray_set_computed():
 
 @pytest.mark.parametrize("test_case", BYTES_SET_TESTCASES)
 def test_StaticBytes_set_py_bytes(test_case: bytes | bytearray):
-    value = abi.StaticBytes(len(test_case))
+    value: abi.StaticBytes = abi.StaticBytes(abi.StaticBytesTypeSpec(len(test_case)))
 
     expr = value.set(test_case)
     assert expr.type_of() == pt.TealType.none
@@ -281,7 +314,9 @@ def test_StaticBytes_set_py_bytes(test_case: bytes | bytearray):
 
 @pytest.mark.parametrize("test_case", BYTES_SET_TESTCASES)
 def test_StaticBytes_expr(test_case: bytes | bytearray):
-    value = abi.StaticBytes(len(test_case) * 2)
+    value: abi.StaticBytes = abi.StaticBytes(
+        abi.StaticBytesTypeSpec(len(test_case) * 2)
+    )
     set_expr = pt.Concat(pt.Bytes(test_case), pt.Bytes(test_case))
 
     expr = value.set(set_expr)

pyteal/ast/abi/util.py

@@ -110,11 +110,13 @@ def type_spec_from_annotation(annotation: Any) -> TypeSpec:
     from pyteal.ast.abi.array_dynamic import (
         DynamicArrayTypeSpec,
         DynamicArray,
+        DynamicBytesTypeSpec,
         DynamicBytes,
     )
     from pyteal.ast.abi.array_static import (
         StaticArrayTypeSpec,
         StaticArray,
+        StaticBytesTypeSpec,
         StaticBytes,
     )
     from pyteal.ast.abi.tuple import (
@@ -220,7 +222,7 @@ def type_spec_from_annotation(annotation: Any) -> TypeSpec:
     if origin is DynamicBytes:
         if len(args) != 0:
             raise TypeError(f"DynamicBytes expect 0 type argument. Got: {args}")
-        return DynamicArrayTypeSpec(ByteTypeSpec())
+        return DynamicBytesTypeSpec()
 
     if origin is DynamicArray:
         if len(args) != 1:
@@ -234,7 +236,7 @@ def type_spec_from_annotation(annotation: Any) -> TypeSpec:
         if len(args) != 1:
             raise TypeError(f"StaticBytes expect 1 type argument. Got: {args}")
         array_length = int_literal_from_annotation(args[0])
-        return StaticArrayTypeSpec(ByteTypeSpec(), array_length)
+        return StaticBytesTypeSpec(array_length)
 
     if origin is StaticArray:
         if len(args) != 2:

pyteal/ast/abi/util_test.py

@@ -277,11 +277,11 @@ class TypeAnnotationTest(NamedTuple):
         ),
         TypeAnnotationTest(
             annotation=abi.StaticBytes[Literal[10]],
-            expected=abi.StaticArrayTypeSpec(abi.ByteTypeSpec(), 10),
+            expected=abi.StaticBytesTypeSpec(10),
         ),
         TypeAnnotationTest(
             annotation=abi.DynamicBytes,
-            expected=abi.DynamicArrayTypeSpec(abi.ByteTypeSpec()),
+            expected=abi.DynamicBytesTypeSpec(),
         ),
     ]