Static/Dynamic-Bytes Convenience (#500)

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Added
 * Add the ability to insert comments in TEAL source file with the Comment method ([#410](https://github.com/algorand/pyteal/pull/410))
+* Static and Dynamic Byte Array convenience classes ([#500](https://github.com/algorand/pyteal/pull/500))
 
 ## Fixed
 * Fix AST duplication bug in `String.set` when called with an `Expr` argument ([#508](https://github.com/algorand/pyteal/pull/508))

docs/abi.rst

@@ -208,7 +208,9 @@ PyTeal Type                                    ARC-4 Type             Dynamic /
 :any:`abi.Byte`                                :code:`byte`           Static                              An 8-bit unsigned integer. This is an alias for :code:`abi.Uint8` that should be used to indicate non-numeric data, such as binary arrays.
 :any:`abi.StaticArray[T,N] <abi.StaticArray>`  :code:`T[N]`           Static when :code:`T` is static     A fixed-length array of :code:`T` with :code:`N` elements
 :any:`abi.Address`                             :code:`address`        Static                              A 32-byte Algorand address. This is an alias for :code:`abi.StaticArray[abi.Byte, Literal[32]]`.
+:any:`abi.StaticBytes[N] <abi.StaticBytes>`    :code:`byte[N]`        Static                              A fixed-length array with :code:`N` elements of :code:`abi.Byte`.
 :any:`abi.DynamicArray[T] <abi.DynamicArray>`  :code:`T[]`            Dynamic                             A variable-length array of :code:`T`
+:any:`abi.DynamicBytes`                        :code:`byte[]`         Dynamic                             A variable-length array of :code:`abi.Byte`.
 :any:`abi.String`                              :code:`string`         Dynamic                             A variable-length byte array assumed to contain UTF-8 encoded content. This is an alias for :code:`abi.DynamicArray[abi.Byte]`.
 :any:`abi.Tuple`\*, :any:`abi.NamedTuple`      :code:`(...)`          Static when all elements are static A tuple of multiple types
 ============================================== ====================== =================================== =======================================================================================================================================================
@@ -263,7 +265,9 @@ All basic types have a :code:`set()` method which can be used to assign a value.
 * :any:`abi.Bool.set(...) <abi.Bool.set>`
 * :any:`abi.StaticArray[T, N].set(...) <abi.StaticArray.set>`
 * :any:`abi.Address.set(...) <abi.Address.set>`
+* :any:`abi.StaticBytes[N].set(...) <abi.StaticBytes.set>`
 * :any:`abi.DynamicArray[T].set(...) <abi.DynamicArray.set>`
+* :any:`abi.DynamicBytes.set(...) <abi.DynamicBytes.set>`
 * :any:`abi.String.set(...) <abi.String.set>`
 * :any:`abi.Tuple.set(...) <abi.Tuple.set>`
 

pyteal/ast/abi/__init__.py

@@ -35,8 +35,12 @@
     Field,
 )
 from pyteal.ast.abi.array_base import ArrayTypeSpec, Array, ArrayElement
-from pyteal.ast.abi.array_static import StaticArrayTypeSpec, StaticArray
-from pyteal.ast.abi.array_dynamic import DynamicArrayTypeSpec, DynamicArray
+from pyteal.ast.abi.array_static import StaticArrayTypeSpec, StaticArray, StaticBytes
+from pyteal.ast.abi.array_dynamic import (
+    DynamicArrayTypeSpec,
+    DynamicArray,
+    DynamicBytes,
+)
 from pyteal.ast.abi.reference_type import (
     ReferenceTypeSpec,
     ReferenceType,
@@ -128,8 +132,10 @@
     "ArrayElement",
     "StaticArrayTypeSpec",
     "StaticArray",
+    "StaticBytes",
     "DynamicArrayTypeSpec",
     "DynamicArray",
+    "DynamicBytes",
     "MethodReturn",
     "Transaction",
     "TransactionTypeSpec",

pyteal/ast/abi/array_dynamic.py

@@ -1,17 +1,13 @@
-from typing import (
-    Union,
-    Sequence,
-    TypeVar,
-    cast,
-)
-
+from typing import Union, Sequence, TypeVar, cast
 
 from pyteal.errors import TealInputError
 from pyteal.ast.expr import Expr
 from pyteal.ast.seq import Seq
+from pyteal.ast.int import Int
+from pyteal.ast.substring import Suffix
 
 from pyteal.ast.abi.type import ComputedValue, BaseType
-from pyteal.ast.abi.uint import Uint16
+from pyteal.ast.abi.uint import Uint16, Byte, ByteTypeSpec
 from pyteal.ast.abi.array_base import ArrayTypeSpec, Array
 
 
@@ -102,3 +98,67 @@ def length(self) -> Expr:
 
 
 DynamicArray.__module__ = "pyteal.abi"
+
+
+class DynamicBytes(DynamicArray[Byte]):
+    """The convenience class that represents ABI dynamic byte array."""
+
+    def __init__(self) -> None:
+        super().__init__(DynamicArrayTypeSpec(ByteTypeSpec()))
+
+    def set(
+        self,
+        values: Union[
+            bytes,
+            bytearray,
+            Expr,
+            Sequence[Byte],
+            DynamicArray[Byte],
+            ComputedValue[DynamicArray[Byte]],
+        ],
+    ) -> Expr:
+        """Set the elements of this DynamicBytes to the input values.
+
+        The behavior of this method depends on the input argument type:
+
+            * :code:`bytes`: set the value to the Python byte string.
+            * :code:`bytearray`: set the value to the Python byte array.
+            * :code:`Expr`: set the value to the result of a PyTeal expression, which must evaluate to a TealType.bytes.
+            * :code:`Sequence[Byte]`: set the bytes of this String to those contained in this Python sequence (e.g. a list or tuple).
+            * :code:`DynamicArray[Byte]`: copy the bytes from another DynamicArray. The argument's element type must exactly match Byte, otherwise an error will occur.
+            * :code:`ComputedValue[DynamicArray[Byte]]`: copy the bytes from a DynamicArray produced by a ComputedValue. The argument's element type must exactly match Byte, otherwise an error will occur.
+
+        Args:
+            values: The new elements this DynamicBytes should have. This must follow the above constraints.
+
+        Returns:
+            An expression which stores the given value into this DynamicBytes.
+        """
+        # NOTE: the import here is to avoid importing in partial initialized module abi
+        from pyteal.ast.abi.string import (
+            _encoded_byte_string,
+            _store_encoded_expr_byte_string_into_var,
+        )
+
+        match values:
+            case bytes() | bytearray():
+                return self.stored_value.store(_encoded_byte_string(values))
+            case Expr():
+                return _store_encoded_expr_byte_string_into_var(
+                    values, self.stored_value
+                )
+
+        return super().set(values)
+
+    def get(self) -> Expr:
+        """Get the byte encoding of this DynamicBytes.
+
+        Dropping the uint16 encoding prefix for dynamic array length.
+
+        Returns:
+            A Pyteal expression that loads byte encoding of this DynamicBytes, and drop the first uint16 DynamicArray length encoding.
+        """
+        return Suffix(self.stored_value.load(), Int(2))
+
+
+DynamicBytes.__module__ = "pyteal.abi"

pyteal/ast/abi/array_dynamic_test.py

@@ -1,7 +1,9 @@
 from typing import List
-import pytest
 
+import pytest
 import pyteal as pt
+
+from algosdk.abi import ABIType
 from pyteal import abi
 from pyteal.ast.abi.util import substring_for_decoding
 from pyteal.ast.abi.tuple import _encode_tuple
@@ -203,6 +205,97 @@ def test_DynamicArray_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),
+]
+
+
+@pytest.mark.parametrize("test_case", BYTES_SET_TESTCASES)
+def test_DynamicBytes_set_py_bytes(test_case: bytes | bytearray):
+    value = abi.DynamicBytes()
+
+    expr = value.set(test_case)
+    assert expr.type_of() == pt.TealType.none
+    assert not expr.has_return()
+
+    actual, _ = expr.__teal__(options)
+    actual.addIncoming()
+    actual = actual.NormalizeBlocks(actual)
+
+    length_encoding = ABIType.from_string("uint16").encode(len(test_case)).hex()
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(None, pt.Op.byte, "0x" + length_encoding + BYTE_HEX_TEST_CASE),
+            pt.TealOp(None, pt.Op.store, value.stored_value.slot),
+        ]
+    )
+
+    with pt.TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+
+@pytest.mark.parametrize("test_case", BYTES_SET_TESTCASES)
+def test_DynamicBytes_set_expr(test_case: bytes | bytearray):
+    value = abi.DynamicBytes()
+
+    set_expr = pt.Concat(pt.Bytes(test_case), pt.Bytes(test_case))
+
+    expr = value.set(set_expr)
+    assert expr.type_of() == pt.TealType.none
+    assert not expr.has_return()
+
+    actual, _ = expr.__teal__(options)
+    actual.addIncoming()
+    actual = actual.NormalizeBlocks(actual)
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(None, pt.Op.byte, "0x" + BYTE_HEX_TEST_CASE),
+            pt.TealOp(None, pt.Op.byte, "0x" + BYTE_HEX_TEST_CASE),
+            pt.TealOp(None, pt.Op.concat),
+            pt.TealOp(None, pt.Op.store, value.stored_value.slot),
+            pt.TealOp(None, pt.Op.load, value.stored_value.slot),
+            pt.TealOp(None, pt.Op.len),
+            pt.TealOp(None, pt.Op.itob),
+            pt.TealOp(None, pt.Op.extract, 6, 0),
+            pt.TealOp(None, pt.Op.load, value.stored_value.slot),
+            pt.TealOp(None, pt.Op.concat),
+            pt.TealOp(None, pt.Op.store, value.stored_value.slot),
+        ]
+    )
+
+    with pt.TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+
+def test_DynamicBytes_get():
+    value = abi.DynamicBytes()
+
+    expr = value.get()
+    assert expr.type_of() == pt.TealType.bytes
+    assert not expr.has_return()
+
+    actual, _ = expr.__teal__(options)
+    actual.addIncoming()
+    actual = actual.NormalizeBlocks(actual)
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(None, pt.Op.load, value.stored_value.slot),
+            pt.TealOp(None, pt.Op.extract, 2, 0),
+        ]
+    )
+
+    with pt.TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+
 def test_DynamicArray_encode():
     dynamicArrayType = abi.DynamicArrayTypeSpec(abi.Uint64TypeSpec())
     value = dynamicArrayType.new_instance()

pyteal/ast/abi/array_static.py

@@ -1,11 +1,16 @@
 from typing import Final, Generic, Literal, Sequence, TypeVar, Union, cast
 
 from pyteal.errors import TealInputError
+from pyteal.ast.assert_ import Assert
 from pyteal.ast.expr import Expr
 from pyteal.ast.int import Int
+from pyteal.ast.bytes import Bytes
+from pyteal.ast.seq import Seq
+from pyteal.ast.unaryexpr import Len
 
 from pyteal.ast.abi.type import ComputedValue, TypeSpec, BaseType
 from pyteal.ast.abi.bool import BoolTypeSpec, _bool_sequence_length
+from pyteal.ast.abi.uint import Byte, ByteTypeSpec
 from pyteal.ast.abi.array_base import ArrayTypeSpec, Array, ArrayElement
 
 
@@ -142,3 +147,64 @@ def __getitem__(self, index: Union[int, Expr]) -> "ArrayElement[T]":
 
 
 StaticArray.__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 set(
+        self,
+        values: Union[
+            bytes,
+            bytearray,
+            Expr,
+            Sequence[Byte],
+            StaticArray[Byte, N],
+            ComputedValue[StaticArray[Byte, N]],
+        ],
+    ) -> Expr:
+        """Set the elements of this StaticBytes to the input values.
+
+        The behavior of this method depends on the input argument type:
+
+            * :code:`bytes`: set the value to the Python byte string.
+            * :code:`bytearray`: set the value to the Python byte array.
+            * :code:`Expr`: set the value to the result of a PyTeal expression, which must evaluate to a TealType.bytes.
+            * :code:`Sequence[Byte]`: set the bytes of this String to those contained in this Python sequence (e.g. a list or tuple).
+            * :code:`StaticArray[Byte, N]`: copy the bytes from another StaticArray. The argument's element type and length must exactly match Byte and this StaticBytes' length, otherwise an error will occur.
+            * :code:`ComputedValue[StaticArray[Byte, N]]`: copy the bytes from a StaticArray produced by a ComputedType. The argument's element type and length must exactly match Byte and this StaticBytes' length, otherwise an error will occur.
+
+        Args:
+            values: The new elements this StaticBytes should have. This must follow the above constraints.
+
+        Returns:
+            An expression which stores the given value into this StaticBytes.
+        """
+        match values:
+            case bytes() | bytearray():
+                if len(values) != self.type_spec().length_static():
+                    raise TealInputError(
+                        f"Got bytes with length {len(values)}, expect {self.type_spec().length_static()}"
+                    )
+                return self.stored_value.store(Bytes(values))
+            case Expr():
+                return Seq(
+                    self.stored_value.store(values),
+                    Assert(self.length() == Len(self.stored_value.load())),
+                )
+
+        return super().set(values)
+
+    def get(self) -> Expr:
+        """Get the byte encoding of this StaticBytes.
+
+        Returns:
+            A Pyteal expression that loads byte encoding of this StaticBytes.
+        """
+        return self.stored_value.load()
+
+
+StaticBytes.__module__ = "pyteal.abi"