From fe41635b4bb7245c86b9975dc11373f2819c3019 Mon Sep 17 00:00:00 2001 From: Hang Su Date: Fri, 1 Jul 2022 11:28:36 -0400 Subject: [PATCH] minor fix --- pyteal/ast/abi/address.py | 9 ++++++++- pyteal/ast/abi/address_test.py | 5 +++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/pyteal/ast/abi/address.py b/pyteal/ast/abi/address.py index 1e8da1f0f..47c559225 100644 --- 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 index eb5ed4681..15408abc8 100644 --- 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), ]