diff --git a/pyteal/ast/abi/__init__.py b/pyteal/ast/abi/__init__.py index e2d57bf29..0dfefb026 100644 --- a/pyteal/ast/abi/__init__.py +++ b/pyteal/ast/abi/__init__.py @@ -1,3 +1,5 @@ +from .string import String, StringTypeSpec +from .address import AddressTypeSpec, Address, ADDRESS_LENGTH from .type import TypeSpec, BaseType, ComputedValue from .bool import BoolTypeSpec, Bool from .uint import ( @@ -32,6 +34,11 @@ from .method_return import MethodReturn __all__ = [ + "String", + "StringTypeSpec", + "Address", + "AddressTypeSpec", + "ADDRESS_LENGTH", "TypeSpec", "BaseType", "ComputedValue", diff --git a/pyteal/ast/abi/address.py b/pyteal/ast/abi/address.py new file mode 100644 index 000000000..0a95ab6f7 --- /dev/null +++ b/pyteal/ast/abi/address.py @@ -0,0 +1,33 @@ +from .array_static import StaticArray, StaticArrayTypeSpec +from .uint import ByteTypeSpec +from ..expr import Expr + +ADDRESS_LENGTH = 32 + + +class AddressTypeSpec(StaticArrayTypeSpec): + def __init__(self) -> None: + super().__init__(ByteTypeSpec(), ADDRESS_LENGTH) + + def new_instance(self) -> "Address": + return Address() + + def __str__(self) -> str: + return "address" + + +AddressTypeSpec.__module__ = "pyteal" + + +class Address(StaticArray): + def __init__(self) -> None: + super().__init__(AddressTypeSpec(), ADDRESS_LENGTH) + + def type_spec(self) -> AddressTypeSpec: + return AddressTypeSpec() + + def get(self) -> Expr: + return self.stored_value.load() + + +Address.__module__ = "pyteal" diff --git a/pyteal/ast/abi/address_test.py b/pyteal/ast/abi/address_test.py new file mode 100644 index 000000000..1227df77b --- /dev/null +++ b/pyteal/ast/abi/address_test.py @@ -0,0 +1,76 @@ +from ... import * + +options = CompileOptions(version=5) + + +def test_AddressTypeSpec_str(): + assert str(abi.AddressTypeSpec()) == "address" + + +def test_AddressTypeSpec_is_dynamic(): + assert (abi.AddressTypeSpec()).is_dynamic() is False + + +def test_AddressTypeSpec_byte_length_static(): + assert (abi.AddressTypeSpec()).byte_length_static() == abi.ADDRESS_LENGTH + + +def test_AddressTypeSpec_new_instance(): + assert isinstance(abi.AddressTypeSpec().new_instance(), abi.Address) + + +def test_AddressTypeSpec_eq(): + assert abi.AddressTypeSpec() == abi.AddressTypeSpec() + + for otherType in ( + abi.ByteTypeSpec, + abi.StaticArrayTypeSpec(abi.ByteTypeSpec(), 31), + abi.DynamicArrayTypeSpec(abi.ByteTypeSpec()), + ): + assert abi.AddressTypeSpec() != otherType + + +def test_Address_encode(): + value = abi.Address() + expr = value.encode() + assert expr.type_of() == TealType.bytes + assert expr.has_return() is False + + expected = TealSimpleBlock([TealOp(expr, Op.load, value.stored_value.slot)]) + actual, _ = expr.__teal__(options) + assert actual == expected + + +def test_Address_decode(): + from os import urandom + + value = abi.Address() + for value_to_set in [urandom(abi.ADDRESS_LENGTH) for x in range(10)]: + expr = value.decode(Bytes(value_to_set)) + + assert expr.type_of() == TealType.none + assert expr.has_return() is False + + expected = TealSimpleBlock( + [ + TealOp(None, Op.byte, f"0x{value_to_set.hex()}"), + TealOp(None, Op.store, value.stored_value.slot), + ] + ) + actual, _ = expr.__teal__(options) + actual.addIncoming() + actual = TealBlock.NormalizeBlocks(actual) + + with TealComponent.Context.ignoreExprEquality(): + assert actual == expected + + +def test_Address_get(): + value = abi.Address() + expr = value.get() + assert expr.type_of() == TealType.bytes + assert expr.has_return() is False + + expected = TealSimpleBlock([TealOp(expr, Op.load, value.stored_value.slot)]) + actual, _ = expr.__teal__(options) + assert actual == expected diff --git a/pyteal/ast/abi/string.py b/pyteal/ast/abi/string.py new file mode 100644 index 000000000..6560bed22 --- /dev/null +++ b/pyteal/ast/abi/string.py @@ -0,0 +1,37 @@ +from .array_dynamic import DynamicArray, DynamicArrayTypeSpec +from .uint import ByteTypeSpec, Uint16TypeSpec +from .util import substringForDecoding + +from ..int import Int +from ..expr import Expr + + +class StringTypeSpec(DynamicArrayTypeSpec): + def __init__(self) -> None: + super().__init__(ByteTypeSpec()) + + def new_instance(self) -> "String": + return String() + + def __str__(self) -> str: + return "string" + + +StringTypeSpec.__module__ = "pyteal" + + +class String(DynamicArray): + def __init__(self) -> None: + super().__init__(StringTypeSpec()) + + def type_spec(self) -> StringTypeSpec: + return StringTypeSpec() + + def get(self) -> Expr: + return substringForDecoding( + self.stored_value.load(), + startIndex=Int(Uint16TypeSpec().byte_length_static()), + ) + + +String.__module__ = "pyteal" diff --git a/pyteal/ast/abi/string_test.py b/pyteal/ast/abi/string_test.py new file mode 100644 index 000000000..856f48ddf --- /dev/null +++ b/pyteal/ast/abi/string_test.py @@ -0,0 +1,79 @@ +from ... import * + +options = CompileOptions(version=5) + + +def test_StringTypeSpec_str(): + assert str(abi.StringTypeSpec()) == "string" + + +def test_StringTypeSpec_is_dynamic(): + assert (abi.StringTypeSpec()).is_dynamic() + + +def test_StringTypeSpec_new_instance(): + assert isinstance(abi.StringTypeSpec().new_instance(), abi.String) + + +def test_StringTypeSpec_eq(): + assert abi.StringTypeSpec() == abi.StringTypeSpec() + + for otherType in ( + abi.ByteTypeSpec, + abi.StaticArrayTypeSpec(abi.ByteTypeSpec(), 1), + abi.DynamicArrayTypeSpec(abi.Uint8TypeSpec()), + ): + assert abi.StringTypeSpec() != otherType + + +def test_String_encode(): + value = abi.String() + expr = value.encode() + assert expr.type_of() == TealType.bytes + assert expr.has_return() is False + + expected = TealSimpleBlock([TealOp(expr, Op.load, value.stored_value.slot)]) + actual, _ = expr.__teal__(options) + assert actual == expected + + +def test_String_decode(): + import random + from os import urandom + + value = abi.String() + for value_to_set in [urandom(random.randint(0, 50)) for x in range(10)]: + expr = value.decode(Bytes(value_to_set)) + + assert expr.type_of() == TealType.none + assert expr.has_return() is False + + expected = TealSimpleBlock( + [ + TealOp(None, Op.byte, f"0x{value_to_set.hex()}"), + TealOp(None, Op.store, value.stored_value.slot), + ] + ) + actual, _ = expr.__teal__(options) + actual.addIncoming() + actual = TealBlock.NormalizeBlocks(actual) + + with TealComponent.Context.ignoreExprEquality(): + assert actual == expected + + +def test_String_get(): + value = abi.String() + expr = value.get() + assert expr.type_of() == TealType.bytes + assert expr.has_return() is False + + expected = TealSimpleBlock( + [TealOp(expr, Op.load, value.stored_value.slot), TealOp(None, Op.extract, 2, 0)] + ) + actual, _ = expr.__teal__(options) + actual.addIncoming() + actual = TealBlock.NormalizeBlocks(actual) + + with TealComponent.Context.ignoreExprEquality(): + assert actual == expected diff --git a/pyteal/ast/abi/util.py b/pyteal/ast/abi/util.py index e9077ded6..f16a03900 100644 --- a/pyteal/ast/abi/util.py +++ b/pyteal/ast/abi/util.py @@ -107,6 +107,8 @@ def type_spec_from_annotation(annotation: Any) -> TypeSpec: Tuple4, Tuple5, ) + from .string import StringTypeSpec, String + from .address import AddressTypeSpec, Address origin = get_origin(annotation) if origin is None: @@ -144,6 +146,16 @@ def type_spec_from_annotation(annotation: Any) -> TypeSpec: raise TypeError("Uint64 expects 0 type arguments. Got: {}".format(args)) return Uint64TypeSpec() + if origin is String: + if len(args) != 0: + raise TypeError("String expects 0 arguments. Got: {}".format(args)) + return StringTypeSpec() + + if origin is Address: + if len(args) != 0: + raise TypeError("Address expects 0 arguments. Got: {}".format(args)) + return AddressTypeSpec() + if origin is DynamicArray: if len(args) != 1: raise TypeError(