Refined ABI type-check on assignment

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Unreleased
 
+## Added
+* Added `abi.type_spec_is_assignable_to` to check for compatible ABI type assignments. ([#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/tuple.py

@@ -15,7 +15,7 @@
 from collections import OrderedDict
 
 from pyteal.types import TealType
-from pyteal.errors import TealInputError
+from pyteal.errors import TealInputError, TealInternalError
 from pyteal.ast.expr import Expr
 from pyteal.ast.seq import Seq
 from pyteal.ast.int import Int
@@ -522,6 +522,13 @@ def __init__(
         self.instance_class: type["NamedTuple"] = instance_class
         super().__init__(*value_type_specs)
 
+    def __eq__(self, other: object) -> bool:
+        return (
+            isinstance(other, NamedTupleTypeSpec)
+            and self.instance_class == other.instance_class
+            and self.value_type_specs() == other.value_type_specs()
+        )
+
     def annotation_type(self) -> "type[NamedTuple]":
         return self.instance_class
 
@@ -552,9 +559,50 @@ class User(abi.NamedTuple):
 
             my_user = User()
 
+    .. automethod:: __init_subclass__
     .. automethod:: __getattr__
     """
 
+    def __init_subclass__(cls) -> None:
+        """This method ensures one only constructs directly from `NamedTuple`,
+        rather than inheriting from `NamedTuple`'s inheritance.
+
+        We want to prohibit the following examples:
+
+            .. code-block:: python
+
+                from pyteal import *
+
+                class LegalInheritance(abi.NamedTuple):
+                    a: abi.Field[abi.Uint64]
+
+                # following are bad cases we guard against
+
+                class IllegalInheritance0(LegalInheritance):
+                    a: abi.Field[abi.Uint64]
+
+                class IllegalInheritance1(LegalInheritance, abi.NamedTuple):
+                    a: abi.Field[abi.Uint64]
+        """
+        is_named_tuple_in_bases = False
+
+        for base_t in cls.__bases__:
+            if base_t == NamedTuple:
+                is_named_tuple_in_bases = True
+            elif issubclass(base_t, NamedTuple):
+                raise TealInternalError(
+                    f"Cannot construct {cls} by inheriting {cls.__bases__}. "
+                    f"Must be constructed by direct inheritance from NamedTuple"
+                )
+
+        if not is_named_tuple_in_bases:
+            raise TealInternalError(
+                "Unexpected: did not find NamedTuple in __bases__,"
+                "did not find NamedTuple in __bases__ member's __bases__"
+            )
+
+        super().__init_subclass__()
+
     def __init__(self):
         if type(self) is NamedTuple:
             raise TealInputError("NamedTuple must be subclassed")

pyteal/ast/abi/tuple_test.py

@@ -925,3 +925,22 @@ def test_NamedTuple_getitem(test_case: type[abi.NamedTuple]):
 
     with pytest.raises(pt.TealInputError):
         tuple_value.f0 = abi.Uint64()
+
+
+def test_NamedTupleTypeSpec():
+    class Point(abi.NamedTuple):
+        x: abi.Field[abi.Uint64]
+        y: abi.Field[abi.Uint64]
+
+    class AccountRecord(abi.NamedTuple):
+        algoBalance: abi.Field[abi.Uint64]
+        assetBalance: abi.Field[abi.Uint64]
+
+    p = Point()
+    ar = AccountRecord()
+
+    assert p.type_spec() == p.type_spec()
+    assert ar.type_spec() == ar.type_spec()
+    assert p.type_spec() != ar.type_spec()
+    assert not abi.type_spec_is_assignable_to(p.type_spec(), ar.type_spec())
+    assert not abi.type_spec_is_assignable_to(ar.type_spec(), p.type_spec())

pyteal/ast/abi/util.py

@@ -499,3 +499,97 @@ 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:
+    """Decides if the value of type :code:`a` can be assigned to or interpreted as another value of type :code:`b`.
+
+    This method return true if and only if all of the following properties hold:
+
+        * value of type :code:`a` has identical encoding as value of type :code:`b`
+        * type :code:`b` is as general as, or more general than type :code:`a`
+
+    For `abi.NamedTuple`, we allow mutual assigning between `abi.Tuple` and `abi.NamedTuple`.
+    But between `abi.NamedTuple`, we only return true when the type specs are identical, or we cannot compare against generality.
+
+    Some examples are illustrated as following:
+
+    =========================== =========================== ============= ====================================================================
+    Type :code:`a`              Type :code:`b`              Assignable?   Reason
+    =========================== =========================== ============= ====================================================================
+    :code:`DynamicArray[Byte]`  :code:`DynamicBytes`        :code:`True`  :code:`DynamicBytes` is as general as :code:`DynamicArray[Byte]`
+    :code:`DynamicBytes`        :code:`DynamicArray[Byte]`  :code:`True`  :code:`DynamicArray[Byte]` is as general as :code:`DynamicBytes`
+    :code:`StaticArray[Byte,N]` :code:`StaticBytes[N]`      :code:`True`  :code:`StaticBytes[N]` is as general as :code:`StaticArray[Byte,N]`
+    :code:`StaticBytes[N]`      :code:`StaticArray[Byte,N]` :code:`True`  :code:`StaticArray[Byte,N]` is as general as :code:`StaticBytes[N]`
+    :code:`String`              :code:`DynamicBytes`        :code:`True`  :code:`DynamicBytes` is more general than :code:`String`
+    :code:`DynamicBytes`        :code:`String`              :code:`False` :code:`String` is more specific than :code:`DynamicBytes`
+    :code:`Address`             :code:`StaticBytes[32]`     :code:`True`  :code:`StaticBytes[32]` is more general than :code:`Address`
+    :code:`StaticBytes[32]`     :code:`Address`             :code:`False` :code:`Address` is more specific than :code:`StaticBytes[32]`
+    :code:`PaymentTransaction`  :code:`Transaction`         :code:`True`  :code:`Transaction` is more general than :code:`PaymentTransaction`
+    :code:`Transaction`         :code:`PaymentTransaction`  :code:`False` :code:`PaymentTransaction` is more specific than :code`Transaction`
+    :code:`Uint8`               :code:`Byte`                :code:`True`  :code:`Uint8` is as general as :code:`Byte`
+    :code:`Byte`                :code:`Uint8`               :code:`True`  :code:`Byte` is as general as :code:`Uint8`
+    =========================== =========================== ============= ====================================================================
+
+    Args:
+        a: The abi.TypeSpec of the value on the right hand side of the assignment.
+        b: The abi.TypeSpec of the value on the left hand side of the assignment.
+
+    Returns:
+        A boolean result on if type :code:`a` is assignable to type :code:`b`.
+    """
+
+    from pyteal.ast.abi import (
+        TupleTypeSpec,
+        NamedTupleTypeSpec,
+        ArrayTypeSpec,
+        StaticArrayTypeSpec,
+        DynamicArrayTypeSpec,
+        StringTypeSpec,
+        AddressTypeSpec,
+        UintTypeSpec,
+    )
+
+    match a, b:
+        case NamedTupleTypeSpec(), NamedTupleTypeSpec():
+            return 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
+        case UintTypeSpec(), UintTypeSpec():
+            a, b = cast(UintTypeSpec, a), cast(UintTypeSpec, b)
+            return a.size == b.size
+
+    if isinstance(a, type(b)):
+        return True
+    elif str(a) == str(b):
+        return True
+
+    return False