diff --git a/pyteal/ast/__init__.py b/pyteal/ast/__init__.py index f0c83d265..d131e7344 100644 --- a/pyteal/ast/__init__.py +++ b/pyteal/ast/__init__.py @@ -140,8 +140,6 @@ # abi import pyteal.ast.abi as abi # noqa: I250 -# from pyteal.ast.abi import * - __all__ = [ "Expr", "LeafExpr", diff --git a/pyteal/ast/abi/__init__.py b/pyteal/ast/abi/__init__.py index ba495f0ac..a48c58d51 100644 --- a/pyteal/ast/abi/__init__.py +++ b/pyteal/ast/abi/__init__.py @@ -1,5 +1,10 @@ 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, + ADDRESS_LENGTH_BYTES, + ADDRESS_LENGTH_STR, +) 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 +43,8 @@ "StringTypeSpec", "Address", "AddressTypeSpec", - "ADDRESS_LENGTH", + "ADDRESS_LENGTH_BYTES", + "ADDRESS_LENGTH_STR", "TypeSpec", "BaseType", "ComputedValue", diff --git a/pyteal/ast/abi/address.py b/pyteal/ast/abi/address.py index 68196d3df..57b82b343 100644 --- a/pyteal/ast/abi/address.py +++ b/pyteal/ast/abi/address.py @@ -2,13 +2,15 @@ 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.expr import Expr -ADDRESS_LENGTH = 32 +ADDRESS_LENGTH_STR = 58 +ADDRESS_LENGTH_BYTES = 32 T = TypeVar("T", bound=BaseType) N = TypeVar("N", bound=int) @@ -16,7 +18,7 @@ class AddressTypeSpec(StaticArrayTypeSpec): def __init__(self) -> None: - super().__init__(ByteTypeSpec(), ADDRESS_LENGTH) + super().__init__(ByteTypeSpec(), ADDRESS_LENGTH_BYTES) def new_instance(self) -> "Address": return Address() @@ -49,6 +51,7 @@ def set( ComputedValue[StaticArray[T, N]], "Address", str, + bytes, Expr, ], ): @@ -60,20 +63,37 @@ def set( ) return value.store_into(self) - if isinstance(value, BaseType): - if not isinstance(value.type_spec(), AddressTypeSpec): - raise TealInputError( - f"Got {value} with type spec {value.type_spec()}, expected AddressTypeSpec" - ) - return self.decode(value.encode()) + elif isinstance(value, BaseType): + + if isinstance(value.type_spec(), AddressTypeSpec): + return self.decode(value.encode()) + + if ( + isinstance(value.type_spec(), StaticArrayTypeSpec) + and isinstance(value.type_spec().value_type_spec(), ByteTypeSpec) + and value.type_spec().length_static() == ADDRESS_LENGTH_BYTES + ): + return self.decode(value.encode()) - if isinstance(value, str): + raise TealInputError( + f"Got {value} with type spec {value.type_spec()}, expected AddressTypeSpec" + ) + + elif isinstance(value, str) and len(value) == ADDRESS_LENGTH_STR: return self.stored_value.store(Addr(value)) - if isinstance(value, Expr): + elif isinstance(value, bytes) and len(value) == ADDRESS_LENGTH_BYTES: + return self.stored_value.store(Bytes(value)) + + elif isinstance(value, Expr): return self.stored_value.store(value) - raise TealInputError(f"Expected str, Address or Expr got {value}") + elif isinstance(value, Sequence): + return super().set(value) + + raise TealInputError( + f"Got {value}, expected str, bytes, Expr, Sequence, Address, or ComputedType" + ) Address.__module__ = "pyteal" diff --git a/pyteal/ast/abi/address_test.py b/pyteal/ast/abi/address_test.py index 373aac860..a8ceb3681 100644 --- a/pyteal/ast/abi/address_test.py +++ b/pyteal/ast/abi/address_test.py @@ -1,6 +1,7 @@ import pytest import pyteal as pt from pyteal import abi +from pyteal.ast.abi.address import ADDRESS_LENGTH_BYTES from pyteal.ast.global_ import Global, GlobalField from pyteal.ast.abi.type_test import ContainerType @@ -19,7 +20,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() == ADDRESS_LENGTH_BYTES def test_AddressTypeSpec_new_instance(): @@ -51,7 +52,7 @@ def test_Address_encode(): def test_Address_decode(): - address = bytes([0] * abi.ADDRESS_LENGTH) + address = bytes([0] * ADDRESS_LENGTH_BYTES) encoded = pt.Bytes(address) for startIndex in (None, pt.Int(0)): @@ -105,7 +106,36 @@ def test_Address_get(): assert actual == expected -def test_Address_set_static(): +def test_Address_set_StaticArray(): + sa = abi.StaticArray( + abi.StaticArrayTypeSpec(abi.ByteTypeSpec(), ADDRESS_LENGTH_BYTES) + ) + for value_to_set in (sa,): + 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.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) @@ -127,7 +157,32 @@ def test_Address_set_static(): assert actual == expected with pytest.raises(pt.TealInputError): - value.set(bytes(32)) + 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_set_expr(): diff --git a/pyteal/ast/abi/string.py b/pyteal/ast/abi/string.py index 9591bd81d..f69aaf28b 100644 --- a/pyteal/ast/abi/string.py +++ b/pyteal/ast/abi/string.py @@ -58,6 +58,7 @@ def set( ComputedValue[DynamicArray[T]], "String", str, + bytes, Expr, ], ) -> Expr: @@ -71,20 +72,31 @@ def set( return value.store_into(self) if isinstance(value, BaseType): - if not isinstance(value.type_spec(), StringTypeSpec): - raise TealInputError( - f"Got {value} with type spec {value.type_spec()}, expected {StringTypeSpec}" - ) - return self.decode(value.encode()) + if isinstance(value.type_spec(), StringTypeSpec): + return self.decode(value.encode()) + + if isinstance(value.type_spec(), DynamicArrayTypeSpec) and isinstance( + value.type_spec().value_type_spec(), ByteTypeSpec + ): + return self.decode(value.encode()) + + raise TealInputError( + f"Got {value} with type spec {value.type_spec()}, expected {StringTypeSpec}" + ) # Assume not length prefixed - if type(value) is str: + if isinstance(value, str) or isinstance(value, bytes): return self.stored_value.store(encoded_string(Bytes(value))) - if not isinstance(value, Expr): - raise TealInputError("Expected Expr, got {}".format(value)) + if isinstance(value, Expr): + return self.stored_value.store(encoded_string(value)) + + if isinstance(value, Sequence): + return super().set(value) - return self.stored_value.store(encoded_string(value)) + raise TealInputError( + f"Got {value}, expected Sequence, DynamicArray, ComputedValue, String, str, bytes, Expr" + ) String.__module__ = "pyteal" diff --git a/pyteal/ast/abi/string_test.py b/pyteal/ast/abi/string_test.py index 80912a56e..a652d2768 100644 --- a/pyteal/ast/abi/string_test.py +++ b/pyteal/ast/abi/string_test.py @@ -108,6 +108,7 @@ def test_String_get(): def test_String_set_static(): + for value_to_set in ("stringy", "😀", "0xDEADBEEF"): value = abi.String() expr = value.set(value_to_set) @@ -133,8 +134,35 @@ def test_String_set_static(): with pt.TealComponent.Context.ignoreExprEquality(): assert actual == expected - with pytest.raises(pt.TealInputError): - value.set(bytes(32)) + 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() + + 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():