Refined ABI type-check on assignment

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Unreleased
 
+## Added
+* Refined ABI type-check on assignment. ([#540](https://github.com/algorand/pyteal/pull/540))
+
 ## Fixed
 * Erroring on constructing an odd length hex string. ([#539](https://github.com/algorand/pyteal/pull/539))
 * Incorrect behavior when overriding a method name ([#550](https://github.com/algorand/pyteal/pull/550))

pyteal/ast/abi/__init__.py

@@ -84,6 +84,7 @@
     type_spec_from_annotation,
     type_specs_from_signature,
     contains_type_spec,
+    type_spec_is_assignable_to,
 )
 
 __all__ = [
@@ -167,4 +168,5 @@
     "algosdk_from_annotation",
     "algosdk_from_type_spec",
     "contains_type_spec",
+    "type_spec_is_assignable_to",
 ]

pyteal/ast/abi/reference_type.py

@@ -1,4 +1,4 @@
-from typing import List, Final, TypeVar, cast
+from typing import Final, TypeVar, cast
 from abc import abstractmethod
 from pyteal.ast.abi.type import BaseType, TypeSpec
 from pyteal.ast.abi.uint import NUM_BITS_IN_BYTE, uint_decode
@@ -215,7 +215,7 @@ def params(self) -> AppParamObject:
 Application.__module__ = "pyteal.abi"
 
 
-ReferenceTypeSpecs: Final[List[TypeSpec]] = [
+ReferenceTypeSpecs: Final[list[TypeSpec]] = [
     AccountTypeSpec(),
     AssetTypeSpec(),
     ApplicationTypeSpec(),

pyteal/ast/abi/transaction.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Union, cast, List, Final
+from typing import Union, cast, Final
 from pyteal.ast.abi.type import BaseType, ComputedValue, TypeSpec
 from pyteal.ast.expr import Expr
 from pyteal.ast.int import Int
@@ -259,7 +259,7 @@ def __init__(self):
 
 ApplicationCallTransaction.__module__ = "pyteal.abi"
 
-TransactionTypeSpecs: Final[List[TypeSpec]] = [
+TransactionTypeSpecs: Final[list[TypeSpec]] = [
     TransactionTypeSpec(),
     PaymentTransactionTypeSpec(),
     KeyRegisterTransactionTypeSpec(),

pyteal/ast/abi/util.py

@@ -499,3 +499,53 @@ def type_specs_from_signature(sig: str) -> tuple[list[TypeSpec], Optional[TypeSp
         return_type = type_spec_from_algosdk(sdk_method.returns.type)
 
     return [type_spec_from_algosdk(arg.type) for arg in sdk_method.args], return_type
+
+
+def type_spec_is_assignable_to(a: TypeSpec, b: TypeSpec) -> bool:
+    from pyteal.ast.abi import (
+        TupleTypeSpec,
+        ArrayTypeSpec,
+        StaticArrayTypeSpec,
+        DynamicArrayTypeSpec,
+        StringTypeSpec,
+        AddressTypeSpec,
+    )
+
+    match a, b:
+        case TupleTypeSpec(), TupleTypeSpec():
+            a, b = cast(TupleTypeSpec, a), cast(TupleTypeSpec, b)
+            if a.length_static() != b.length_static():
+                return False
+            return all(
+                map(
+                    lambda ab: type_spec_is_assignable_to(ab[0], ab[1]),
+                    zip(a.value_type_specs(), b.value_type_specs()),
+                )
+            )
+        case ArrayTypeSpec(), ArrayTypeSpec():
+            a, b = cast(ArrayTypeSpec, a), cast(ArrayTypeSpec, b)
+            if not type_spec_is_assignable_to(a.value_type_spec(), b.value_type_spec()):
+                return False
+            match a, b:
+                case AddressTypeSpec(), StaticArrayTypeSpec():
+                    a, b = cast(AddressTypeSpec, a), cast(StaticArrayTypeSpec, b)
+                    return a.length_static() == b.length_static()
+                case StaticArrayTypeSpec(), AddressTypeSpec():
+                    return False
+                case StaticArrayTypeSpec(), StaticArrayTypeSpec():
+                    a, b = cast(StaticArrayTypeSpec, a), cast(StaticArrayTypeSpec, b)
+                    return a.length_static() == b.length_static()
+                case StringTypeSpec(), DynamicArrayTypeSpec():
+                    return True
+                case DynamicArrayTypeSpec(), StringTypeSpec():
+                    return False
+                case DynamicArrayTypeSpec(), DynamicArrayTypeSpec():
+                    return True
+            return False
+
+    if isinstance(a, type(b)):
+        return True
+    elif str(a) == str(b):
+        return True
+
+    return False

pyteal/ast/abi/util_test.py

@@ -772,3 +772,250 @@ def test_sdk_type_specs_from_signature(sig_str, sig_args, sig_rets):
     args, ret = abi.type_specs_from_signature(sig_str)
     assert args == sig_args
     assert ret == sig_rets
+
+
+class NamedTDecl(abi.NamedTuple):
+    a: abi.Field[abi.Uint64]
+    b: abi.Field[
+        abi.Tuple3[abi.PaymentTransaction, abi.Address, abi.StaticBytes[Literal[16]]]
+    ]
+    c: abi.Field[abi.Transaction]
+
+
+class NamedTComp0(abi.NamedTuple):
+    a0: abi.Field[abi.String]
+    a1: abi.Field[abi.Address]
+
+
+class NamedTComp1(abi.NamedTuple):
+    b0: abi.Field[abi.DynamicBytes]
+    b1: abi.Field[abi.StaticBytes[Literal[32]]]
+
+
+class NamedTComp2(abi.NamedTuple):
+    b1: abi.Field[abi.Address]
+    b0: abi.Field[abi.DynamicBytes]
+
+
+class SafeBidirectional(NamedTuple):
+    xs: list[abi.TypeSpec]
+
+
+SAFE_BIDIRECTIONAL_TEST_CASES: list[SafeBidirectional] = (
+    [
+        SafeBidirectional(
+            [
+                abi.type_spec_from_annotation(abi.StaticArray[abi.Byte, Literal[10]]),
+                abi.type_spec_from_annotation(abi.StaticBytes[Literal[10]]),
+            ],
+        ),
+        SafeBidirectional([abi.BoolTypeSpec(), abi.BoolTypeSpec()]),
+        SafeBidirectional([abi.StringTypeSpec(), abi.StringTypeSpec()]),
+        SafeBidirectional([abi.AddressTypeSpec(), abi.AddressTypeSpec()]),
+        SafeBidirectional(
+            [
+                abi.type_spec_from_annotation(abi.DynamicBytes),
+                abi.type_spec_from_annotation(abi.DynamicArray[abi.Byte]),
+            ]
+        ),
+        SafeBidirectional(
+            [
+                abi.type_spec_from_annotation(
+                    abi.Tuple3[
+                        abi.Uint64,
+                        abi.Tuple3[
+                            abi.PaymentTransaction,
+                            abi.Address,
+                            abi.StaticArray[abi.Byte, Literal[16]],
+                        ],
+                        abi.Transaction,
+                    ]
+                ),
+                abi.type_spec_from_annotation(NamedTDecl),
+                abi.type_spec_from_annotation(NamedTDecl),
+            ]
+        ),
+    ]
+    + [
+        SafeBidirectional([spec, spec])
+        for spec in abi.ReferenceTypeSpecs + abi.TransactionTypeSpecs
+    ]
+    + [
+        SafeBidirectional([spec_t(), spec_t()])
+        for spec_t in bfs_on_inheritance(abi.UintTypeSpec)
+        if not isabstract(spec_t)
+    ]
+)
+
+
+@pytest.mark.parametrize("tc", SAFE_BIDIRECTIONAL_TEST_CASES)
+def test_type_spec_is_assignable_safe_bidirectional(tc: SafeBidirectional):
+    assert len(tc.xs) > 0
+    for a in tc.xs:
+        for b in tc.xs:
+            assert abi.type_spec_is_assignable_to(a, b)
+
+
+@pytest.mark.parametrize("ts", bfs_on_inheritance(abi.TypeSpec))
+def test_type_spec_is_assignable_safe_bidirectional_full_coverage(ts: type):
+    def exists_in_safe_bidirectional(_ts: type):
+        for safe_bidirectional in SAFE_BIDIRECTIONAL_TEST_CASES:
+            for t in safe_bidirectional.xs:
+                if type(t) == _ts:
+                    return True
+        return False
+
+    if isabstract(ts):
+        return
+    assert exists_in_safe_bidirectional(ts)
+
+
+class SafeAssignment(NamedTuple):
+    a: abi.TypeSpec
+    bs: list[abi.TypeSpec]
+
+
+SAFE_ASSIGNMENT_TEST_CASES: list[SafeAssignment] = [
+    SafeAssignment(
+        abi.StringTypeSpec(),
+        [abi.DynamicBytesTypeSpec(), abi.DynamicArrayTypeSpec(abi.ByteTypeSpec())],
+    ),
+    SafeAssignment(
+        abi.type_spec_from_annotation(NamedTDecl),
+        [
+            abi.type_spec_from_annotation(
+                abi.Tuple3[
+                    abi.Uint64,
+                    abi.Tuple3[
+                        abi.Transaction,
+                        abi.Address,
+                        abi.StaticArray[abi.Byte, Literal[16]],
+                    ],
+                    abi.Transaction,
+                ]
+            ),
+        ],
+    ),
+    SafeAssignment(
+        abi.type_spec_from_annotation(
+            abi.Tuple3[
+                abi.Uint64,
+                abi.Tuple3[
+                    abi.PaymentTransaction,
+                    abi.Address,
+                    abi.StaticArray[abi.Byte, Literal[16]],
+                ],
+                abi.PaymentTransaction,
+            ]
+        ),
+        [abi.type_spec_from_annotation(NamedTDecl)],
+    ),
+    SafeAssignment(
+        abi.AddressTypeSpec(),
+        [abi.StaticArrayTypeSpec(abi.ByteTypeSpec(), 32), abi.StaticBytesTypeSpec(32)],
+    ),
+    SafeAssignment(
+        abi.type_spec_from_annotation(NamedTComp0),
+        [abi.type_spec_from_annotation(NamedTComp1)],
+    ),
+] + [
+    SafeAssignment(spec, [abi.TransactionTypeSpec()])
+    for spec in abi.TransactionTypeSpecs
+    if spec != abi.TransactionTypeSpec()
+]
+
+
+@pytest.mark.parametrize("tc", SAFE_ASSIGNMENT_TEST_CASES)
+def test_type_spec_is_assignable_safe_assignment(tc: SafeAssignment):
+    assert len(tc.bs) > 0
+    for b in tc.bs:
+        assert abi.type_spec_is_assignable_to(tc.a, b)
+        assert not abi.type_spec_is_assignable_to(b, tc.a)
+
+
+@pytest.mark.parametrize("ts", bfs_on_inheritance(abi.TypeSpec))
+def test_type_spec_is_assignable_safe_assignment_full_coverage(ts: type):
+    def exists_in_safe_assignment(_ts: type):
+        for safe_assignment in SAFE_ASSIGNMENT_TEST_CASES:
+            if type(safe_assignment.a) == _ts:
+                return True
+            for t in safe_assignment.bs:
+                if type(t) == _ts:
+                    return True
+        return False
+
+    if isabstract(ts) or all(isabstract(parent) for parent in ts.__bases__):
+        return
+    assert exists_in_safe_assignment(ts)
+
+
+class UnsafeBidirectional(NamedTuple):
+    xs: list[abi.TypeSpec]
+
+
+UNSAFE_BIDIRECTIONAL_TEST_CASES: list[UnsafeBidirectional] = [
+    UnsafeBidirectional(
+        [abi.DynamicArrayTypeSpec(abi.Uint8TypeSpec()), abi.DynamicBytesTypeSpec()]
+    ),
+    UnsafeBidirectional(
+        [abi.StaticBytesTypeSpec(7), abi.StaticArrayTypeSpec(abi.ByteTypeSpec(), 11)]
+    ),
+    UnsafeBidirectional(
+        [
+            abi.StringTypeSpec(),
+            abi.DynamicArrayTypeSpec(abi.Uint32TypeSpec()),
+            abi.DynamicArrayTypeSpec(abi.Uint64TypeSpec()),
+            abi.AddressTypeSpec(),
+            abi.StaticBytesTypeSpec(33),
+        ]
+    ),
+    UnsafeBidirectional(
+        [
+            abi.type_spec_from_annotation(NamedTDecl),
+            abi.type_spec_from_annotation(NamedTComp0),
+            abi.type_spec_from_annotation(NamedTComp2),
+        ]
+    ),
+    UnsafeBidirectional(
+        [
+            abi.type_spec_from_annotation(NamedTComp1),
+            abi.TupleTypeSpec(
+                abi.AddressTypeSpec(),
+                abi.StaticArrayTypeSpec(abi.Uint16TypeSpec(), 100),
+            ),
+        ]
+    ),
+    UnsafeBidirectional(
+        [spec for spec in abi.TransactionTypeSpecs if spec != abi.TransactionTypeSpec()]
+        + [
+            spec()
+            for spec in bfs_on_inheritance(abi.UintTypeSpec)
+            if not isabstract(spec)
+        ]
+        + [abi.BoolTypeSpec()]
+    ),
+    UnsafeBidirectional(abi.ReferenceTypeSpecs + [abi.TransactionTypeSpec()]),
+]
+
+
+@pytest.mark.parametrize("tc", UNSAFE_BIDIRECTIONAL_TEST_CASES)
+def test_type_spec_is_assignable_unsafe_bidirectional(tc: UnsafeBidirectional):
+    for ia, a in enumerate(tc.xs):
+        for ib, b in enumerate(tc.xs):
+            if ia == ib:
+                continue
+            assert not abi.type_spec_is_assignable_to(a, b)
+
+
+@pytest.mark.parametrize("ts", bfs_on_inheritance(abi.TypeSpec))
+def test_type_spec_is_assignable_unsafe_bidirectional_full_coverage(ts: type):
+    def exists_in_unsafe_bidirectional(_ts: type):
+        for unsafe_bidirectional in UNSAFE_BIDIRECTIONAL_TEST_CASES:
+            for t in unsafe_bidirectional.xs:
+                if type(t) == _ts:
+                    return True
+        return False
+
+    if isabstract(ts):
+        return
+    assert exists_in_unsafe_bidirectional(ts)