adding handlers for other cases, adding testing for new handlers

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

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

pyteal/ast/abi/address.py

@@ -2,21 +2,23 @@
 
 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)
 
 
 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"

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():

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"

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():