Add set on address and string, change array init to accept typespec

pyteal/ast/abi/__init__.py

@@ -1,5 +1,9 @@
 from pyteal.ast.abi.string import String, StringTypeSpec
-from pyteal.ast.abi.address import AddressTypeSpec, Address, ADDRESS_LENGTH
+from pyteal.ast.abi.address import (
+    AddressTypeSpec,
+    Address,
+    AddressLength,
+)
 from pyteal.ast.abi.type import TypeSpec, BaseType, ComputedValue
 from pyteal.ast.abi.bool import BoolTypeSpec, Bool
 from pyteal.ast.abi.uint import (
@@ -38,7 +42,7 @@
     "StringTypeSpec",
     "Address",
     "AddressTypeSpec",
-    "ADDRESS_LENGTH",
+    "AddressLength",
     "TypeSpec",
     "BaseType",
     "ComputedValue",

pyteal/ast/abi/address.py

@@ -1,33 +1,104 @@
+from enum import IntEnum
+from typing import Union, Sequence, Literal, cast
+from collections.abc import Sequence as CollectionSequence
+
+from pyteal.errors import TealInputError
+
+from pyteal.ast.bytes import Bytes
+from pyteal.ast.addr import Addr
+from pyteal.ast.abi.type import ComputedValue, BaseType
 from pyteal.ast.abi.array_static import StaticArray, StaticArrayTypeSpec
-from pyteal.ast.abi.uint import ByteTypeSpec
+from pyteal.ast.abi.uint import ByteTypeSpec, Byte
 from pyteal.ast.expr import Expr
 
-ADDRESS_LENGTH = 32
+
+class AddressLength(IntEnum):
+    String = 58
+    Bytes = 32
+
+
+AddressLength.__module__ = "pyteal"
 
 
 class AddressTypeSpec(StaticArrayTypeSpec):
     def __init__(self) -> None:
-        super().__init__(ByteTypeSpec(), ADDRESS_LENGTH)
+        super().__init__(ByteTypeSpec(), AddressLength.Bytes)
 
     def new_instance(self) -> "Address":
         return Address()
 
     def __str__(self) -> str:
         return "address"
 
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, AddressTypeSpec)
+
 
 AddressTypeSpec.__module__ = "pyteal"
 
 
-class Address(StaticArray):
+class Address(StaticArray[Byte, Literal[AddressLength.Bytes]]):
     def __init__(self) -> None:
-        super().__init__(AddressTypeSpec(), ADDRESS_LENGTH)
+        super().__init__(AddressTypeSpec())
 
     def type_spec(self) -> AddressTypeSpec:
         return AddressTypeSpec()
 
     def get(self) -> Expr:
         return self.stored_value.load()
 
+    def set(
+        self,
+        value: Union[
+            Sequence[Byte],
+            StaticArray[Byte, Literal[AddressLength.Bytes]],
+            ComputedValue[StaticArray[Byte, Literal[AddressLength.Bytes]]],
+            "Address",
+            str,
+            bytes,
+            Expr,
+        ],
+    ):
+
+        match value:
+            case ComputedValue():
+                pts = value.produced_type_spec()
+                if pts == AddressTypeSpec() or pts == StaticArrayTypeSpec(
+                    ByteTypeSpec(), AddressLength.Bytes
+                ):
+                    return value.store_into(self)
+
+                raise TealInputError(
+                    f"Got ComputedValue with type spec {pts}, expected AddressTypeSpec or StaticArray[Byte, Literal[AddressLength.Bytes]]"
+                )
+            case BaseType():
+                if (
+                    value.type_spec() == AddressTypeSpec()
+                    or value.type_spec()
+                    == StaticArrayTypeSpec(ByteTypeSpec(), AddressLength.Bytes)
+                ):
+                    return self.stored_value.store(value.stored_value.load())
+
+                raise TealInputError(
+                    f"Got {value} with type spec {value.type_spec()}, expected AddressTypeSpec"
+                )
+            case str():
+                # Addr throws if value is invalid address
+                return self.stored_value.store(Addr(value))
+            case bytes():
+                if len(value) == AddressLength.Bytes:
+                    return self.stored_value.store(Bytes(value))
+                raise TealInputError(
+                    f"Got bytes with length {len(value)}, expected {AddressLength.Bytes}"
+                )
+            case Expr():
+                return self.stored_value.store(value)
+            case CollectionSequence():
+                return super().set(cast(Sequence[Byte], value))
+
+        raise TealInputError(
+            f"Got {type(value)}, expected Sequence, StaticArray, ComputedValue, Address, str, bytes, Expr"
+        )
+
 
 Address.__module__ = "pyteal"

pyteal/ast/abi/address_test.py

@@ -1,6 +1,11 @@
+import pytest
 import pyteal as pt
 from pyteal import abi
 
+from pyteal.ast.abi.type_test import ContainerType
+from pyteal.ast.abi.util import substringForDecoding
+
+
 options = pt.CompileOptions(version=5)
 
 
@@ -13,7 +18,7 @@ def test_AddressTypeSpec_is_dynamic():
 
 
 def test_AddressTypeSpec_byte_length_static():
-    assert (abi.AddressTypeSpec()).byte_length_static() == abi.ADDRESS_LENGTH
+    assert (abi.AddressTypeSpec()).byte_length_static() == abi.AddressLength.Bytes
 
 
 def test_AddressTypeSpec_new_instance():
@@ -24,8 +29,8 @@ def test_AddressTypeSpec_eq():
     assert abi.AddressTypeSpec() == abi.AddressTypeSpec()
 
     for otherType in (
-        abi.ByteTypeSpec,
-        abi.StaticArrayTypeSpec(abi.ByteTypeSpec(), 31),
+        abi.ByteTypeSpec(),
+        abi.StaticArrayTypeSpec(abi.ByteTypeSpec(), 32),
         abi.DynamicArrayTypeSpec(abi.ByteTypeSpec()),
     ):
         assert abi.AddressTypeSpec() != otherType
@@ -45,37 +50,212 @@ def test_Address_encode():
 
 
 def test_Address_decode():
-    from os import urandom
+    address = bytes([0] * abi.AddressLength.Bytes)
+    encoded = pt.Bytes(address)
+
+    for startIndex in (None, pt.Int(0)):
+        for endIndex in (None, pt.Int(1)):
+            for length in (None, pt.Int(2)):
+                value = abi.Address()
+
+                if endIndex is not None and length is not None:
+                    with pytest.raises(pt.TealInputError):
+                        value.decode(
+                            encoded,
+                            startIndex=startIndex,
+                            endIndex=endIndex,
+                            length=length,
+                        )
+                    continue
+
+                expr = value.decode(
+                    encoded, startIndex=startIndex, endIndex=endIndex, length=length
+                )
+                assert expr.type_of() == pt.TealType.none
+                assert expr.has_return() is False
+
+                expectedExpr = value.stored_value.store(
+                    substringForDecoding(
+                        encoded, startIndex=startIndex, endIndex=endIndex, length=length
+                    )
+                )
+                expected, _ = expectedExpr.__teal__(options)
+                expected.addIncoming()
+                expected = pt.TealBlock.NormalizeBlocks(expected)
+
+                actual, _ = expr.__teal__(options)
+                actual.addIncoming()
+                actual = pt.TealBlock.NormalizeBlocks(actual)
+
+                with pt.TealComponent.Context.ignoreExprEquality():
+                    assert actual == expected
+
+
+def test_Address_get():
+    value = abi.Address()
+    expr = value.get()
+    assert expr.type_of() == pt.TealType.bytes
+    assert expr.has_return() is False
+
+    expected = pt.TealSimpleBlock(
+        [pt.TealOp(expr, pt.Op.load, value.stored_value.slot)]
+    )
+    actual, _ = expr.__teal__(options)
+    assert actual == expected
+
 
+def test_Address_set_StaticArray():
+    value_to_set = abi.StaticArray(
+        abi.StaticArrayTypeSpec(abi.ByteTypeSpec(), abi.AddressLength.Bytes)
+    )
     value = abi.Address()
-    for value_to_set in [urandom(abi.ADDRESS_LENGTH) for x in range(10)]:
-        expr = value.decode(pt.Bytes(value_to_set))
+    expr = value.set(value_to_set)
+    assert expr.type_of() == pt.TealType.none
+    assert not expr.has_return()
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(None, pt.Op.load, value_to_set.stored_value.slot),
+            pt.TealOp(None, pt.Op.store, value.stored_value.slot),
+        ]
+    )
+
+    actual, _ = expr.__teal__(options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
 
+    with pt.TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+    with pytest.raises(pt.TealInputError):
+        bogus = abi.StaticArray(abi.StaticArrayTypeSpec(abi.ByteTypeSpec(), 10))
+        value.set(bogus)
+
+
+def test_Address_set_str():
+    for value_to_set in ("CEZZTYHNTVIZFZWT6X2R474Z2P3Q2DAZAKIRTPBAHL3LZ7W4O6VBROVRQA",):
+        value = abi.Address()
+        expr = value.set(value_to_set)
         assert expr.type_of() == pt.TealType.none
-        assert expr.has_return() is False
+        assert not expr.has_return()
+
+        expected = pt.TealSimpleBlock(
+            [
+                pt.TealOp(None, pt.Op.addr, value_to_set),
+                pt.TealOp(None, pt.Op.store, value.stored_value.slot),
+            ]
+        )
+
+        actual, _ = expr.__teal__(options)
+        actual.addIncoming()
+        actual = pt.TealBlock.NormalizeBlocks(actual)
+
+        with pt.TealComponent.Context.ignoreExprEquality():
+            assert actual == expected
+
+        with pytest.raises(pt.TealInputError):
+            value.set(" " * 16)
+
+
+def test_Address_set_bytes():
+    for value_to_set in (bytes(32),):
+        value = abi.Address()
+        expr = value.set(value_to_set)
+        assert expr.type_of() == pt.TealType.none
+        assert not expr.has_return()
 
         expected = pt.TealSimpleBlock(
             [
                 pt.TealOp(None, pt.Op.byte, f"0x{value_to_set.hex()}"),
                 pt.TealOp(None, pt.Op.store, value.stored_value.slot),
             ]
         )
+
         actual, _ = expr.__teal__(options)
         actual.addIncoming()
         actual = pt.TealBlock.NormalizeBlocks(actual)
 
         with pt.TealComponent.Context.ignoreExprEquality():
             assert actual == expected
 
+        with pytest.raises(pt.TealInputError):
+            value.set(bytes(16))
 
-def test_Address_get():
+        with pytest.raises(pt.TealInputError):
+            value.set(16)
+
+
+def test_Address_set_expr():
+    for value_to_set in [pt.Global(pt.GlobalField.zero_address)]:
+        value = abi.Address()
+        expr = value.set(value_to_set)
+        assert expr.type_of() == pt.TealType.none
+        assert not expr.has_return()
+
+        vts, _ = value_to_set.__teal__(options)
+        expected = pt.TealSimpleBlock(
+            [
+                vts.ops[0],
+                pt.TealOp(None, pt.Op.store, value.stored_value.slot),
+            ]
+        )
+
+        actual, _ = expr.__teal__(options)
+        actual.addIncoming()
+        actual = pt.TealBlock.NormalizeBlocks(actual)
+
+        with pt.TealComponent.Context.ignoreExprEquality():
+            assert actual == expected
+
+
+def test_Address_set_copy():
     value = abi.Address()
-    expr = value.get()
-    assert expr.type_of() == pt.TealType.bytes
-    assert expr.has_return() is False
+    other = abi.Address()
+    expr = value.set(other)
+    assert expr.type_of() == pt.TealType.none
+    assert not expr.has_return()
 
     expected = pt.TealSimpleBlock(
-        [pt.TealOp(expr, pt.Op.load, value.stored_value.slot)]
+        [
+            pt.TealOp(None, pt.Op.load, other.stored_value.slot),
+            pt.TealOp(None, pt.Op.store, value.stored_value.slot),
+        ]
     )
+
     actual, _ = expr.__teal__(options)
-    assert actual == expected
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    with pt.TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+    with pytest.raises(pt.TealInputError):
+        value.set(abi.String())
+
+
+def test_Address_set_computed():
+    av = pt.Addr("MDDKJUCTY57KA2PBFI44CLTJ5YHY5YVS4SVQUPZAWSRV2ZAVFKI33O6YPE")
+    computed_value = ContainerType(abi.AddressTypeSpec(), av)
+
+    value = abi.Address()
+    expr = value.set(computed_value)
+    assert expr.type_of() == pt.TealType.none
+    assert not expr.has_return()
+
+    _, byte_ops = av.__teal__(options)
+    expected = pt.TealSimpleBlock(
+        [
+            byte_ops.ops[0],
+            pt.TealOp(None, pt.Op.store, value.stored_value.slot),
+        ]
+    )
+
+    actual, _ = expr.__teal__(options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    with pt.TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+    with pytest.raises(pt.TealInputError):
+        value.set(ContainerType(abi.ByteTypeSpec(), pt.Int(0x01)))

pyteal/ast/abi/array_base_test.py

@@ -107,7 +107,7 @@ def test_ArrayElement_store_into():
             assert actual == expected
 
         with pytest.raises(pt.TealInputError):
-            element.store_into(abi.Tuple(elementType))
+            element.store_into(abi.Tuple(abi.TupleTypeSpec(elementType)))
 
     for elementType in STATIC_TYPES + DYNAMIC_TYPES:
         dynamicArrayType = abi.DynamicArrayTypeSpec(elementType)
@@ -158,4 +158,4 @@ def test_ArrayElement_store_into():
         )
 
         with pytest.raises(pt.TealInputError):
-            element.store_into(abi.Tuple(elementType))
+            element.store_into(abi.Tuple(abi.TupleTypeSpec(elementType)))