minor fix

algorand · Jul 1, 2022 · fe41635 · fe41635
1 parent 8de49ab
commit fe41635
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/pyteal/ast/abi/address.py b/pyteal/ast/abi/address.py
@@ -4,7 +4,11 @@
 
 from pyteal.errors import TealInputError
 
+from pyteal.ast.assert_ import Assert
 from pyteal.ast.bytes import Bytes
+from pyteal.ast.int import Int
+from pyteal.ast.seq import Seq
+from pyteal.ast.unaryexpr import Len
 from pyteal.ast.addr import Addr
 from pyteal.ast.abi.type import ComputedValue, BaseType
 from pyteal.ast.abi.array_static import StaticArray, StaticArrayTypeSpec
@@ -95,7 +99,10 @@ def set(
                     f"Got bytes with length {len(value)}, expected {AddressLength.Bytes}"
                 )
             case Expr():
-                return self.stored_value.store(value)
+                return Seq(
+                    Assert(Len(value) == Int(AddressLength.Bytes.value)),
+                    self.stored_value.store(value),
+                )
             case CollectionSequence():
                 return super().set(cast(Sequence[Byte], value))
 

diff --git a/pyteal/ast/abi/address_test.py b/pyteal/ast/abi/address_test.py
@@ -199,6 +199,11 @@ def test_Address_set_expr():
         vts, _ = value_to_set.__teal__(options)
         expected = pt.TealSimpleBlock(
             [
+                vts.ops[0],
+                pt.TealOp(None, pt.Op.len),
+                pt.TealOp(None, pt.Op.int, 32),
+                pt.TealOp(None, pt.Op.eq),
+                pt.TealOp(None, pt.Op.assert_),
                 vts.ops[0],
                 pt.TealOp(None, pt.Op.store, value.stored_value.slot),
             ]