diff --git a/pyteal/ast/abi/__init__.py b/pyteal/ast/abi/__init__.py index ba495f0ac..89629e612 100644 --- a/pyteal/ast/abi/__init__.py +++ b/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", diff --git a/pyteal/ast/abi/address.py b/pyteal/ast/abi/address.py index 1e6a44095..db21b1adc 100644 --- a/pyteal/ast/abi/address.py +++ b/pyteal/ast/abi/address.py @@ -1,13 +1,28 @@ +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() @@ -15,13 +30,16 @@ def new_instance(self) -> "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() @@ -29,5 +47,58 @@ def type_spec(self) -> 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" diff --git a/pyteal/ast/abi/address_test.py b/pyteal/ast/abi/address_test.py index b6d5bde39..c3ce39510 100644 --- a/pyteal/ast/abi/address_test.py +++ b/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,14 +50,119 @@ 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( [ @@ -60,6 +170,7 @@ def test_Address_decode(): pt.TealOp(None, pt.Op.store, value.stored_value.slot), ] ) + actual, _ = expr.__teal__(options) actual.addIncoming() actual = pt.TealBlock.NormalizeBlocks(actual) @@ -67,15 +178,84 @@ def test_Address_decode(): 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))) diff --git a/pyteal/ast/abi/array_base_test.py b/pyteal/ast/abi/array_base_test.py index 0c480183a..7b515cbc6 100644 --- a/pyteal/ast/abi/array_base_test.py +++ b/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))) diff --git a/pyteal/ast/abi/array_dynamic.py b/pyteal/ast/abi/array_dynamic.py index f515b4ac8..99d0ffffe 100644 --- a/pyteal/ast/abi/array_dynamic.py +++ b/pyteal/ast/abi/array_dynamic.py @@ -10,7 +10,7 @@ from pyteal.ast.expr import Expr from pyteal.ast.seq import Seq -from pyteal.ast.abi.type import ComputedValue, TypeSpec, BaseType +from pyteal.ast.abi.type import ComputedValue, BaseType from pyteal.ast.abi.uint import Uint16 from pyteal.ast.abi.array_base import ArrayTypeSpec, Array @@ -20,7 +20,7 @@ class DynamicArrayTypeSpec(ArrayTypeSpec[T]): def new_instance(self) -> "DynamicArray[T]": - return DynamicArray(self.value_type_spec()) + return DynamicArray(self) def is_length_dynamic(self) -> bool: return True @@ -47,8 +47,8 @@ def __str__(self) -> str: class DynamicArray(Array[T]): """The class that represents ABI dynamic array type.""" - def __init__(self, value_type_spec: TypeSpec) -> None: - super().__init__(DynamicArrayTypeSpec(value_type_spec)) + def __init__(self, array_type_spec: DynamicArrayTypeSpec[T]) -> None: + super().__init__(array_type_spec) def type_spec(self) -> DynamicArrayTypeSpec[T]: return cast(DynamicArrayTypeSpec[T], super().type_spec()) diff --git a/pyteal/ast/abi/array_dynamic_test.py b/pyteal/ast/abi/array_dynamic_test.py index 6095fba3b..8d44de478 100644 --- a/pyteal/ast/abi/array_dynamic_test.py +++ b/pyteal/ast/abi/array_dynamic_test.py @@ -170,7 +170,7 @@ def test_DynamicArray_set_copy(): def test_DynamicArray_set_computed(): - value = abi.DynamicArray(abi.ByteTypeSpec()) + value = abi.DynamicArray(abi.DynamicArrayTypeSpec(abi.ByteTypeSpec())) computed = ContainerType( value.type_spec(), pt.Bytes("this should be a dynamic array") ) diff --git a/pyteal/ast/abi/array_static.py b/pyteal/ast/abi/array_static.py index f9e8f834e..6f8ba4468 100644 --- a/pyteal/ast/abi/array_static.py +++ b/pyteal/ast/abi/array_static.py @@ -28,7 +28,7 @@ def __init__(self, value_type_spec: TypeSpec, array_length: int) -> None: self.array_length: Final = array_length def new_instance(self) -> "StaticArray[T, N]": - return StaticArray(self.value_type_spec(), self.length_static()) + return StaticArray(self) def length_static(self) -> int: """Get the size of this static array type. @@ -72,8 +72,8 @@ def __str__(self) -> str: class StaticArray(Array[T], Generic[T, N]): """The class that represents ABI static array type.""" - def __init__(self, value_type_spec: TypeSpec, array_length: int) -> None: - super().__init__(StaticArrayTypeSpec(value_type_spec, array_length)) + def __init__(self, array_type_spec: StaticArrayTypeSpec[T, N]) -> None: + super().__init__(array_type_spec) def type_spec(self) -> StaticArrayTypeSpec[T, N]: return cast(StaticArrayTypeSpec[T, N], super().type_spec()) diff --git a/pyteal/ast/abi/array_static_test.py b/pyteal/ast/abi/array_static_test.py index 70cc817fc..826cef0e4 100644 --- a/pyteal/ast/abi/array_static_test.py +++ b/pyteal/ast/abi/array_static_test.py @@ -104,7 +104,9 @@ def test_StaticArray_decode(): for startIndex in (None, pt.Int(1)): for endIndex in (None, pt.Int(2)): for length in (None, pt.Int(3)): - value = abi.StaticArray(abi.Uint64TypeSpec(), 10) + value = abi.StaticArray( + abi.StaticArrayTypeSpec(abi.Uint64TypeSpec(), 10) + ) if endIndex is not None and length is not None: with pytest.raises(pt.TealInputError): @@ -140,7 +142,7 @@ def test_StaticArray_decode(): def test_StaticArray_set_values(): - value = abi.StaticArray(abi.Uint64TypeSpec(), 10) + value = abi.StaticArray(abi.StaticArrayTypeSpec(abi.Uint64TypeSpec(), 10)) with pytest.raises(pt.TealInputError): value.set([]) @@ -176,14 +178,14 @@ def test_StaticArray_set_values(): def test_StaticArray_set_copy(): - value = abi.StaticArray(abi.Uint64TypeSpec(), 10) - otherArray = abi.StaticArray(abi.Uint64TypeSpec(), 10) + value = abi.StaticArray(abi.StaticArrayTypeSpec(abi.Uint64TypeSpec(), 10)) + otherArray = abi.StaticArray(abi.StaticArrayTypeSpec(abi.Uint64TypeSpec(), 10)) with pytest.raises(pt.TealInputError): - value.set(abi.StaticArray(abi.Uint64TypeSpec(), 11)) + value.set(abi.StaticArray(abi.StaticArrayTypeSpec(abi.Uint64TypeSpec(), 11))) with pytest.raises(pt.TealInputError): - value.set(abi.StaticArray(abi.Uint8TypeSpec(), 10)) + value.set(abi.StaticArray(abi.StaticArrayTypeSpec(abi.Uint8TypeSpec(), 10))) with pytest.raises(pt.TealInputError): value.set(abi.Uint64()) @@ -208,7 +210,7 @@ def test_StaticArray_set_copy(): def test_StaticArray_set_computed(): - value = abi.StaticArray(abi.Uint64TypeSpec(), 10) + value = abi.StaticArray(abi.StaticArrayTypeSpec(abi.Uint64TypeSpec(), 10)) computed = ContainerType( value.type_spec(), pt.Bytes("indeed this is hard to simulate") ) @@ -239,7 +241,7 @@ def test_StaticArray_set_computed(): def test_StaticArray_encode(): - value = abi.StaticArray(abi.Uint64TypeSpec(), 10) + value = abi.StaticArray(abi.StaticArrayTypeSpec(abi.Uint64TypeSpec(), 10)) expr = value.encode() assert expr.type_of() == pt.TealType.bytes assert not expr.has_return() @@ -258,7 +260,7 @@ def test_StaticArray_encode(): def test_StaticArray_length(): for length in (0, 1, 2, 3, 1000): - value = abi.StaticArray(abi.Uint64TypeSpec(), length) + value = abi.StaticArray(abi.StaticArrayTypeSpec(abi.Uint64TypeSpec(), length)) expr = value.length() assert expr.type_of() == pt.TealType.uint64 assert not expr.has_return() @@ -275,7 +277,7 @@ def test_StaticArray_length(): def test_StaticArray_getitem(): for length in (0, 1, 2, 3, 1000): - value = abi.StaticArray(abi.Uint64TypeSpec(), length) + value = abi.StaticArray(abi.StaticArrayTypeSpec(abi.Uint64TypeSpec(), length)) for index in range(length): # dynamic indexes diff --git a/pyteal/ast/abi/method_return_test.py b/pyteal/ast/abi/method_return_test.py index d834ddf2a..00657b6be 100644 --- a/pyteal/ast/abi/method_return_test.py +++ b/pyteal/ast/abi/method_return_test.py @@ -7,7 +7,7 @@ POSITIVE_CASES = [ abi.Uint16(), abi.Uint32(), - abi.StaticArray(abi.BoolTypeSpec(), 12), + abi.StaticArray(abi.StaticArrayTypeSpec(abi.BoolTypeSpec(), 12)), ] diff --git a/pyteal/ast/abi/string.py b/pyteal/ast/abi/string.py index ca32024a1..cddaa4e7d 100644 --- a/pyteal/ast/abi/string.py +++ b/pyteal/ast/abi/string.py @@ -1,9 +1,25 @@ +from typing import Union, TypeVar, Sequence +from collections.abc import Sequence as CollectionSequence + +from pyteal.ast.abi.type import ComputedValue, BaseType from pyteal.ast.abi.array_dynamic import DynamicArray, DynamicArrayTypeSpec from pyteal.ast.abi.uint import ByteTypeSpec, Uint16TypeSpec -from pyteal.ast.abi.util import substringForDecoding from pyteal.ast.int import Int from pyteal.ast.expr import Expr +from pyteal.ast.bytes import Bytes +from pyteal.ast.unaryexpr import Itob, Len +from pyteal.ast.substring import Suffix +from pyteal.ast.naryexpr import Concat + +from pyteal.errors import TealInputError + + +def encoded_string(s: Expr): + return Concat(Suffix(Itob(Len(s)), Int(6)), s) + + +T = TypeVar("T", bound=BaseType) class StringTypeSpec(DynamicArrayTypeSpec): @@ -16,6 +32,9 @@ def new_instance(self) -> "String": def __str__(self) -> str: return "string" + def __eq__(self, other: object) -> bool: + return isinstance(other, StringTypeSpec) + StringTypeSpec.__module__ = "pyteal" @@ -28,9 +47,49 @@ def type_spec(self) -> StringTypeSpec: return StringTypeSpec() def get(self) -> Expr: - return substringForDecoding( - self.stored_value.load(), - startIndex=Int(Uint16TypeSpec().byte_length_static()), + return Suffix( + self.stored_value.load(), Int(Uint16TypeSpec().byte_length_static()) + ) + + def set( + self, + value: Union[ + Sequence[T], + DynamicArray[T], + ComputedValue[DynamicArray[T]], + "String", + str, + bytes, + Expr, + ], + ) -> Expr: + + match value: + case ComputedValue(): + if value.produced_type_spec() == StringTypeSpec(): + return value.store_into(self) + + raise TealInputError( + f"Got ComputedValue with type spec {value.produced_type_spec()}, expected StringTypeSpec" + ) + case BaseType(): + if value.type_spec() == StringTypeSpec() or ( + value.type_spec() == DynamicArrayTypeSpec(ByteTypeSpec()) + ): + return self.stored_value.store(value.stored_value.load()) + + raise TealInputError( + f"Got {value} with type spec {value.type_spec()}, expected {StringTypeSpec}" + ) + case str() | bytes(): + return self.stored_value.store(encoded_string(Bytes(value))) + case Expr(): + return self.stored_value.store(encoded_string(value)) + case CollectionSequence(): + return super().set(value) + + raise TealInputError( + f"Got {type(value)}, expected DynamicArray, ComputedValue, String, str, bytes, Expr" ) diff --git a/pyteal/ast/abi/string_test.py b/pyteal/ast/abi/string_test.py index c722c0d7a..a652d2768 100644 --- a/pyteal/ast/abi/string_test.py +++ b/pyteal/ast/abi/string_test.py @@ -1,5 +1,10 @@ +import pytest + import pyteal as pt from pyteal import abi +from pyteal.ast.abi.util import substringForDecoding +from pyteal.ast.abi.type_test import ContainerType +from pyteal.util import escapeStr options = pt.CompileOptions(version=5) @@ -20,9 +25,10 @@ def test_StringTypeSpec_eq(): assert abi.StringTypeSpec() == abi.StringTypeSpec() for otherType in ( - abi.ByteTypeSpec, + abi.ByteTypeSpec(), abi.StaticArrayTypeSpec(abi.ByteTypeSpec(), 1), abi.DynamicArrayTypeSpec(abi.Uint8TypeSpec()), + abi.DynamicArrayTypeSpec(abi.ByteTypeSpec()), ): assert abi.StringTypeSpec() != otherType @@ -40,23 +46,87 @@ def test_String_encode(): assert actual == expected -def test_String_decode(): - import random - from os import urandom +def test_DynamicArray_decode(): + encoded = pt.Bytes("encoded") + stringType = abi.StringTypeSpec() + for startIndex in (None, pt.Int(1)): + for endIndex in (None, pt.Int(2)): + for length in (None, pt.Int(3)): + value = stringType.new_instance() + + 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_String_get(): value = abi.String() - for value_to_set in [urandom(random.randint(0, 50)) for x in range(10)]: - expr = value.decode(pt.Bytes(value_to_set)) + 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), + pt.TealOp(None, pt.Op.extract, 2, 0), + ] + ) + actual, _ = expr.__teal__(options) + actual.addIncoming() + actual = pt.TealBlock.NormalizeBlocks(actual) + + with pt.TealComponent.Context.ignoreExprEquality(): + assert actual == expected + +def test_String_set_static(): + + for value_to_set in ("stringy", "😀", "0xDEADBEEF"): + value = abi.String() + 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.byte, f"0x{value_to_set.hex()}"), + pt.TealOp(None, pt.Op.byte, escapeStr(value_to_set)), + 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.byte, escapeStr(value_to_set)), + pt.TealOp(None, pt.Op.concat), pt.TealOp(None, pt.Op.store, value.stored_value.slot), ] ) + actual, _ = expr.__teal__(options) actual.addIncoming() actual = pt.TealBlock.NormalizeBlocks(actual) @@ -64,22 +134,113 @@ def test_String_decode(): with pt.TealComponent.Context.ignoreExprEquality(): assert actual == expected + for value_to_set in (bytes(32), b"alphabet_soup"): + value = abi.String() + expr = value.set(value_to_set) + assert expr.type_of() == pt.TealType.none + assert not expr.has_return() -def test_String_get(): + teal_val = f"0x{value_to_set.hex()}" + + expected = pt.TealSimpleBlock( + [ + pt.TealOp(None, pt.Op.byte, teal_val), + 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.byte, teal_val), + pt.TealOp(None, pt.Op.concat), + 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(42) + + +def test_String_set_expr(): + for value_to_set in (pt.Bytes("hi"), pt.Bytes("base16", "0xdeadbeef")): + value = abi.String() + 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.len), + pt.TealOp(None, pt.Op.itob), + pt.TealOp(None, pt.Op.extract, 6, 0), + vts.ops[0], + pt.TealOp(None, pt.Op.concat), + 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_String_set_copy(): value = abi.String() - expr = value.get() - assert expr.type_of() == pt.TealType.bytes - assert expr.has_return() is False + other = abi.String() + 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.extract, 2, 0), + pt.TealOp(None, pt.Op.load, other.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): + value.set(abi.Address()) + + +def test_String_set_computed(): + bv = pt.Bytes("base16", "0x0004DEADBEEF") + computed_value = ContainerType(abi.StringTypeSpec(), bv) + + value = abi.String() + expr = value.set(computed_value) + assert expr.type_of() == pt.TealType.none + assert not expr.has_return() + + _, byte_ops = bv.__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))) diff --git a/pyteal/ast/abi/tuple.py b/pyteal/ast/abi/tuple.py index 6d5792829..3824e701e 100644 --- a/pyteal/ast/abi/tuple.py +++ b/pyteal/ast/abi/tuple.py @@ -218,7 +218,7 @@ def length_static(self) -> int: return len(self.value_specs) def new_instance(self) -> "Tuple": - return Tuple(*self.value_specs) + return Tuple(self) def is_dynamic(self) -> bool: return any(type_spec.is_dynamic() for type_spec in self.value_type_specs()) @@ -248,8 +248,8 @@ def __str__(self) -> str: class Tuple(BaseType): - def __init__(self, *value_type_specs: TypeSpec) -> None: - super().__init__(TupleTypeSpec(*value_type_specs)) + def __init__(self, tuple_type_spec: TupleTypeSpec) -> None: + super().__init__(tuple_type_spec) def type_spec(self) -> TupleTypeSpec: return cast(TupleTypeSpec, super().type_spec()) @@ -301,7 +301,7 @@ def length(self) -> Expr: def __getitem__(self, index: int) -> "TupleElement": if not (0 <= index < self.type_spec().length_static()): - raise TealInputError("Index out of bounds") + raise TealInputError(f"Index out of bounds: {index}") return TupleElement(self, index) @@ -336,11 +336,18 @@ def store_into(self, output: BaseType) -> Expr: # sizes. +def _tuple_raise_arg_mismatch(expected: int, typespec: TupleTypeSpec): + if len(typespec.value_specs) != expected: + raise TealInputError( + f"Expected TupleTypeSpec with {expected} elements, Got {len(typespec.value_specs)}" + ) + + class Tuple0(Tuple): """A Tuple with 0 values.""" def __init__(self) -> None: - super().__init__() + super().__init__(TupleTypeSpec()) Tuple0.__module__ = "pyteal" @@ -351,8 +358,9 @@ def __init__(self) -> None: class Tuple1(Tuple, Generic[T1]): """A Tuple with 1 value.""" - def __init__(self, value1_type_spec: TypeSpec) -> None: - super().__init__(value1_type_spec) + def __init__(self, value_type_spec: TupleTypeSpec) -> None: + _tuple_raise_arg_mismatch(1, value_type_spec) + super().__init__(value_type_spec) Tuple1.__module__ = "pyteal" @@ -363,8 +371,9 @@ def __init__(self, value1_type_spec: TypeSpec) -> None: class Tuple2(Tuple, Generic[T1, T2]): """A Tuple with 2 values.""" - def __init__(self, value1_type_spec: TypeSpec, value2_type_spec: TypeSpec) -> None: - super().__init__(value1_type_spec, value2_type_spec) + def __init__(self, value_type_spec: TupleTypeSpec) -> None: + _tuple_raise_arg_mismatch(2, value_type_spec) + super().__init__(value_type_spec) Tuple2.__module__ = "pyteal" @@ -377,11 +386,10 @@ class Tuple3(Tuple, Generic[T1, T2, T3]): def __init__( self, - value1_type_spec: TypeSpec, - value2_type_spec: TypeSpec, - value3_type_spec: TypeSpec, + value_type_spec: TupleTypeSpec, ) -> None: - super().__init__(value1_type_spec, value2_type_spec, value3_type_spec) + _tuple_raise_arg_mismatch(3, value_type_spec) + super().__init__(value_type_spec) Tuple3.__module__ = "pyteal" @@ -394,14 +402,10 @@ class Tuple4(Tuple, Generic[T1, T2, T3, T4]): def __init__( self, - value1_type_spec: TypeSpec, - value2_type_spec: TypeSpec, - value3_type_spec: TypeSpec, - value4_type_spec: TypeSpec, + value_type_spec: TupleTypeSpec, ) -> None: - super().__init__( - value1_type_spec, value2_type_spec, value3_type_spec, value4_type_spec - ) + _tuple_raise_arg_mismatch(4, value_type_spec) + super().__init__(value_type_spec) Tuple4.__module__ = "pyteal" @@ -414,19 +418,10 @@ class Tuple5(Tuple, Generic[T1, T2, T3, T4, T5]): def __init__( self, - value1_type_spec: TypeSpec, - value2_type_spec: TypeSpec, - value3_type_spec: TypeSpec, - value4_type_spec: TypeSpec, - value5_type_spec: TypeSpec, + value_type_spec: TupleTypeSpec, ) -> None: - super().__init__( - value1_type_spec, - value2_type_spec, - value3_type_spec, - value4_type_spec, - value5_type_spec, - ) + _tuple_raise_arg_mismatch(5, value_type_spec) + super().__init__(value_type_spec) Tuple5.__module__ = "pyteal" diff --git a/pyteal/ast/abi/tuple_test.py b/pyteal/ast/abi/tuple_test.py index 6fb9401e4..dbee4e444 100644 --- a/pyteal/ast/abi/tuple_test.py +++ b/pyteal/ast/abi/tuple_test.py @@ -23,10 +23,10 @@ class EncodeTest(NamedTuple): uint16_b = abi.Uint16() bool_a = abi.Bool() bool_b = abi.Bool() - tuple_a = abi.Tuple(abi.BoolTypeSpec(), abi.BoolTypeSpec()) - dynamic_array_a = abi.DynamicArray(abi.Uint64TypeSpec()) - dynamic_array_b = abi.DynamicArray(abi.Uint16TypeSpec()) - dynamic_array_c = abi.DynamicArray(abi.BoolTypeSpec()) + tuple_a = abi.Tuple(abi.TupleTypeSpec(abi.BoolTypeSpec(), abi.BoolTypeSpec())) + dynamic_array_a = abi.DynamicArray(abi.DynamicArrayTypeSpec(abi.Uint64TypeSpec())) + dynamic_array_b = abi.DynamicArray(abi.DynamicArrayTypeSpec(abi.Uint16TypeSpec())) + dynamic_array_c = abi.DynamicArray(abi.DynamicArrayTypeSpec(abi.BoolTypeSpec())) tail_holder = pt.ScratchVar() encoded_tail = pt.ScratchVar() @@ -599,7 +599,7 @@ def test_TupleTypeSpec_byte_length_static(): def test_Tuple_decode(): encoded = pt.Bytes("encoded") - tupleValue = abi.Tuple(abi.Uint64TypeSpec()) + tupleValue = abi.Tuple(abi.TupleTypeSpec(abi.Uint64TypeSpec())) for startIndex in (None, pt.Int(1)): for endIndex in (None, pt.Int(2)): for length in (None, pt.Int(3)): @@ -638,7 +638,9 @@ def test_Tuple_decode(): def test_Tuple_set(): tupleValue = abi.Tuple( - abi.Uint8TypeSpec(), abi.Uint16TypeSpec(), abi.Uint32TypeSpec() + abi.TupleTypeSpec( + abi.Uint8TypeSpec(), abi.Uint16TypeSpec(), abi.Uint32TypeSpec() + ) ) uint8 = abi.Uint8() uint16 = abi.Uint16() @@ -678,7 +680,9 @@ def test_Tuple_set(): def test_Tuple_set_Computed(): tupleValue = abi.Tuple( - abi.Uint8TypeSpec(), abi.Uint16TypeSpec(), abi.Uint32TypeSpec() + abi.TupleTypeSpec( + abi.Uint8TypeSpec(), abi.Uint16TypeSpec(), abi.Uint32TypeSpec() + ) ) computed = ContainerType( tupleValue.type_spec(), pt.Bytes("internal representation") @@ -710,7 +714,7 @@ def test_Tuple_set_Computed(): def test_Tuple_encode(): - tupleValue = abi.Tuple(abi.Uint64TypeSpec()) + tupleValue = abi.Tuple(abi.TupleTypeSpec(abi.Uint64TypeSpec())) expr = tupleValue.encode() assert expr.type_of() == pt.TealType.bytes assert not expr.has_return() @@ -739,7 +743,7 @@ def test_Tuple_length(): ] for i, test in enumerate(tests): - tupleValue = abi.Tuple(*test) + tupleValue = abi.Tuple(abi.TupleTypeSpec(*test)) expr = tupleValue.length() assert expr.type_of() == pt.TealType.uint64 assert not expr.has_return() @@ -767,7 +771,7 @@ def test_Tuple_getitem(): ] for i, test in enumerate(tests): - tupleValue = abi.Tuple(*test) + tupleValue = abi.Tuple(abi.TupleTypeSpec(*test)) for j in range(len(test)): element = tupleValue[j] assert type(element) is TupleElement, "Test at index {} failed".format(i) @@ -793,7 +797,7 @@ def test_TupleElement_store_into(): ] for i, test in enumerate(tests): - tupleValue = abi.Tuple(*test) + tupleValue = abi.Tuple(abi.TupleTypeSpec(*test)) for j in range(len(test)): element = TupleElement(tupleValue, j) output = test[j].new_instance() diff --git a/pyteal/ast/subroutine_test.py b/pyteal/ast/subroutine_test.py index 9f1f91c7b..ee6ae24f2 100644 --- a/pyteal/ast/subroutine_test.py +++ b/pyteal/ast/subroutine_test.py @@ -110,8 +110,12 @@ def fnWithMixedAnns4(a: pt.ScratchVar, b, c: pt.abi.Uint16) -> pt.Expr: x = pt.Int(42) s = pt.Bytes("hello") av_u16 = pt.abi.Uint16() - av_bool_dym_arr = pt.abi.DynamicArray(pt.abi.BoolTypeSpec()) - av_u32_static_arr = pt.abi.StaticArray(pt.abi.Uint32TypeSpec(), 10) + av_bool_dym_arr = pt.abi.DynamicArray( + pt.abi.DynamicArrayTypeSpec(pt.abi.BoolTypeSpec()) + ) + av_u32_static_arr = pt.abi.StaticArray( + pt.abi.StaticArrayTypeSpec(pt.abi.Uint32TypeSpec(), 10) + ) av_bool = pt.abi.Bool() av_byte = pt.abi.Byte() @@ -265,7 +269,9 @@ def fnWithMixedAnns4AndBytesReturn(a: pt.Expr, b: pt.ScratchVar) -> pt.Bytes: def fnWithMixedAnnsABIRet1( a: pt.Expr, b: pt.ScratchVar, c: pt.abi.Uint16 ) -> pt.abi.StaticArray[pt.abi.Uint32, Literal[10]]: - return pt.abi.StaticArray(pt.abi.Uint32TypeSpec(), 10) + return pt.abi.StaticArray( + pt.abi.StaticArrayTypeSpec(pt.abi.Uint32TypeSpec(), 10) + ) def fnWithMixedAnnsABIRet2( a: pt.Expr, b: pt.abi.Byte, c: pt.ScratchVar