ABI Strings

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",

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"

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

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"

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

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(