NamedTuple Implementation (#473)

Co-authored-by: Jason Paulos <jasonpaulos@users.noreply.github.com>
Co-authored-by: Zeph Grunschlag <tzaffi@users.noreply.github.com>
Co-authored-by: Jason Paulos <jasonpaulos@users.noreply.github.com>

docs/abi.rst

@@ -210,7 +210,7 @@ PyTeal Type                                    ARC-4 Type             Dynamic /
 :any:`abi.Address`                             :code:`address`        Static                              A 32-byte Algorand address. This is an alias for :code:`abi.StaticArray[abi.Byte, Literal[32]]`.
 :any:`abi.DynamicArray[T] <abi.DynamicArray>`  :code:`T[]`            Dynamic                             A variable-length array of :code:`T`
 :any:`abi.String`                              :code:`string`         Dynamic                             A variable-length byte array assumed to contain UTF-8 encoded content. This is an alias for :code:`abi.DynamicArray[abi.Byte]`.
-:any:`abi.Tuple`\*                             :code:`(...)`          Static when all elements are static A tuple of multiple types
+:any:`abi.Tuple`\*, :any:`abi.NamedTuple`      :code:`(...)`          Static when all elements are static A tuple of multiple types
 ============================================== ====================== =================================== =======================================================================================================================================================
 
 .. note::
@@ -223,6 +223,21 @@ PyTeal Type                                    ARC-4 Type             Dynamic /
     * :any:`abi.Tuple4[T1,T2,T3,T4] <abi.Tuple4>`: a tuple of four values, :code:`(T1,T2,T3,T4)`
     * :any:`abi.Tuple5[T1,T2,T3,T4,T5] <abi.Tuple5>`: a tuple of five values, :code:`(T1,T2,T3,T4,T5)`
 
+    While we are still on PyTeal 3.10, we have a workaround for :any:`abi.Tuple` by :any:`abi.NamedTuple`, which allows one to define a tuple with more than 5 generic arguments, and access tuple elements by field name. For example:
+
+    .. code-block:: python
+
+        from pyteal import *
+        from typing import Literal as L
+
+        class InheritedFromNamedTuple(abi.NamedTuple):
+            acct_address: abi.Field[abi.Address]
+            amount: abi.Field[abi.Uint64]
+            retrivable: abi.Field[abi.Bool]
+            desc: abi.Field[abi.String]
+            list_of_addrs: abi.Field[abi.DynamicArray[abi.Address]]
+            balance_list: abi.Field[abi.StaticArray[abi.Uint64, L[10]]]
+
 These ARC-4 types are not yet supported in PyTeal:
 
 * Non-power-of-2 unsigned integers under 64 bits, i.e. :code:`uint24`, :code:`uint48`, :code:`uint56`
@@ -285,7 +300,7 @@ All basic types that represent a single value have a :code:`get()` method, which
 A brief example is below. Please consult the documentation linked above for each method to learn more about specific usage and behavior.
 
 .. code-block:: python
-    
+
     from pyteal import *
 
     @Subroutine(TealType.uint64)
@@ -306,11 +321,17 @@ The supported methods are:
 
 * :any:`abi.StaticArray.__getitem__(index: int | Expr) <abi.StaticArray.__getitem__>`, used for :any:`abi.StaticArray` and :any:`abi.Address`
 * :any:`abi.Array.__getitem__(index: int | Expr) <abi.Array.__getitem__>`, used for :any:`abi.DynamicArray` and :any:`abi.String`
-* :any:`abi.Tuple.__getitem__(index: int) <abi.Tuple.__getitem__>`
+* :any:`abi.Tuple.__getitem__(index: int) <abi.Tuple.__getitem__>`, used for :any:`abi.Tuple` and :any:`abi.NamedTuple`\*
 
 .. note::
     Be aware that these methods return a :any:`ComputedValue`, similar to other PyTeal operations which return ABI types. More information about why that is necessary and how to use a :any:`ComputedValue` can be found in the :ref:`Computed Values` section.
 
+.. note::
+    \*For :any:`abi.NamedTuple`, one can access tuple elements through both methods
+
+    * :any:`abi.Tuple.__getitem__(index: int) <abi.Tuple.__getitem__>`
+    * :any:`abi.NamedTuple.__getattr__(name: str) <abi.NamedTuple.__getattr__>`
+
 A brief example is below. Please consult the documentation linked above for each method to learn more about specific usage and behavior.
 
 .. code-block:: python
@@ -759,7 +780,7 @@ Registering Methods
     **We strongly recommend** methods immediately access and validate compound type parameters *before* persisting arguments for later transactions. For validation, it is sufficient to attempt to extract each element your method will use. If there is an input error for an element, indexing into that element will fail.
 
     Notes:
-    
+
     * This recommendation applies to recursively contained compound types as well. Successfully extracting an element which is a compound type does not guarantee the extracted value is valid; you must also inspect its elements as well.
     * Because of this, :any:`abi.Address` is **not** guaranteed to have exactly 32 bytes. To defend against unintended behavior, manually verify the length is 32 bytes, i.e. :code:`Assert(Len(address.get()) == Int(32))`.
 
@@ -804,7 +825,7 @@ The first way to register a method is with the :any:`Router.add_method_handler`
             # store the result in the sender's local state too
             App.localPut(Txn.sender(), Bytes("result", output.get())),
         )
-    
+
     # Register the `add` method with the router, using the default `MethodConfig`
     # (only no-op, non-creation calls allowed).
     router.add_method_handler(add)

pyteal/ast/abi/__init__.py

@@ -30,6 +30,9 @@
     Tuple3,
     Tuple4,
     Tuple5,
+    NamedTuple,
+    NamedTupleTypeSpec,
+    Field,
 )
 from pyteal.ast.abi.array_base import ArrayTypeSpec, Array, ArrayElement
 from pyteal.ast.abi.array_static import StaticArrayTypeSpec, StaticArray
@@ -117,6 +120,9 @@
     "Tuple3",
     "Tuple4",
     "Tuple5",
+    "NamedTuple",
+    "NamedTupleTypeSpec",
+    "Field",
     "ArrayTypeSpec",
     "Array",
     "ArrayElement",

pyteal/ast/abi/tuple.py

@@ -1,3 +1,4 @@
+from inspect import get_annotations
 from typing import (
     List,
     Sequence,
@@ -7,7 +8,11 @@
     cast,
     overload,
     Any,
+    TypeAlias,
+    get_args,
+    get_origin,
 )
+from collections import OrderedDict
 
 from pyteal.types import TealType
 from pyteal.errors import TealInputError
@@ -31,7 +36,7 @@
     _bool_aware_static_byte_length,
 )
 from pyteal.ast.abi.uint import NUM_BITS_IN_BYTE, Uint16
-from pyteal.ast.abi.util import substring_for_decoding
+from pyteal.ast.abi.util import substring_for_decoding, type_spec_from_annotation
 
 
 def _encode_tuple(values: Sequence[BaseType]) -> Expr:
@@ -484,3 +489,135 @@ def __init__(
 
 
 Tuple5.__module__ = "pyteal.abi"
+
+
+Field: TypeAlias = TupleElement[T]
+
+
+class NamedTupleTypeSpec(TupleTypeSpec):
+    """A NamedTupleType inherits from TupleTypeSpec, allowing for more than 5 elements."""
+
+    def __init__(
+        self, instance_class: type["NamedTuple"], *value_type_specs: TypeSpec
+    ) -> None:
+        if instance_class == NamedTuple:
+            raise TealInputError(
+                "NamedTupleTypeSpec must be instanced with subclassed NamedTuple class."
+            )
+
+        self.instance_class: type["NamedTuple"] = instance_class
+        super().__init__(*value_type_specs)
+
+    def annotation_type(self) -> "type[NamedTuple]":
+        return self.instance_class
+
+    def new_instance(self) -> "NamedTuple":
+        return self.instance_class()
+
+
+NamedTupleTypeSpec.__module__ = "pyteal.abi"
+
+
+class NamedTuple(Tuple):
+    """A NamedTuple is a :any:`Tuple` that has named elements, inspired by Python's `typing.NamedTuple <https://docs.python.org/3/library/typing.html#typing.NamedTuple>`_.
+
+    A new NamedTuple type can be created by subclassing this class and adding field annotations.
+    Every field annotation must be an instantiable ABI type wrapped in the :code:`abi.Field` annotation.
+
+    For example:
+
+        .. code-block:: python
+
+            from pyteal import *
+
+            class User(abi.NamedTuple):
+                address: abi.Field[abi.Address]
+                balance: abi.Field[abi.Uint64]
+
+            # User is equivalent to abi.Tuple2[abi.Address, abi.Uint64]
+
+            my_user = User()
+
+    .. automethod:: __getattr__
+    """
+
+    def __init__(self):
+        if type(self) is NamedTuple:
+            raise TealInputError("NamedTuple must be subclassed")
+
+        anns = get_annotations(type(self))
+        if not anns:
+            raise TealInputError("Expected fields to be declared but found none")
+
+        # NOTE: this `_ready` variable enables `__setattr__` during `__init__` execution,
+        # while after `__init__`, we cannot use `__setattr__` to set fields in `NamedTuple`.
+        # NOTE: If we declare variable `__ready`, then internally,
+        # the variable name would be changed to `_NamedTuple__ready`, which is implicit.
+        self.__ready = False
+        self.__type_specs: OrderedDict[str, TypeSpec] = OrderedDict()
+        self.__field_index: dict[str, int] = {}
+
+        for index, (name, annotation) in enumerate(anns.items()):
+            origin = get_origin(annotation)
+            if origin is None:
+                origin = annotation
+            if origin is not get_origin(Field):
+                raise TealInputError(
+                    f'Type annotation for attribute "{name}" must be a Field. Got {origin}'
+                )
+
+            args = get_args(annotation)
+            if len(args) != 1:
+                raise TealInputError(
+                    f'Type annotation for attribute "{name}" must have a single argument. Got {args}'
+                )
+
+            self.__type_specs[name] = type_spec_from_annotation(args[0])
+            self.__field_index[name] = index
+
+        super().__init__(
+            NamedTupleTypeSpec(type(self), *list(self.__type_specs.values()))
+        )
+
+        self.__ready = True
+
+    def __getattr__(self, field: str) -> TupleElement[Any]:
+        """Retrieve an element by its field in this NamedTuple.
+
+        For example:
+
+        .. code-block:: python
+
+            from pyteal import *
+
+            class User(abi.NamedTuple):
+                address: abi.Field[abi.Address]
+                balance: abi.Field[abi.Uint64]
+
+            @ABIReturnSubroutine
+            def get_user_balance(user: User, *, output: abi.Uint64) -> Expr:
+                return output.set(user.balance)
+
+        Args:
+            field: a Python string containing the field to access.
+                This function will raise an KeyError if the field is not available in the defined NamedTuple.
+
+        Returns:
+            A TupleElement that corresponds to the element at the given field name, returning a ComputedType.
+            Due to Python type limitations, the parameterized type of the TupleElement is Any.
+        """
+        return self.__getitem__(self.__field_index[field])
+
+    def __setattr__(self, name: str, field: Any) -> None:
+        # we allow `__setattr__` only when:
+        # - we are in `__init__`: `not self.__ready`
+        # - we are setting `_ready`: `name == "_NamedTuple__ready"`,
+        #   it is internally changed from `__ready` to `_NamedTuple__ready`,
+        #   see notes in `__init__`
+        if name == "_NamedTuple__ready" or not self.__ready:
+            super().__setattr__(name, field)
+            return
+        raise TealInputError("cannot assign to NamedTuple attributes.")
+
+
+NamedTuple.__module__ = "pyteal.abi"

pyteal/ast/abi/tuple_test.py

@@ -1,4 +1,4 @@
-from typing import NamedTuple, List, Callable
+from typing import NamedTuple, List, Callable, Literal
 import pytest
 
 import pyteal as pt
@@ -824,3 +824,104 @@ def test_TupleElement_store_into():
 
             with pt.TealComponent.Context.ignoreExprEquality():
                 assert actual == expected, "Test at index {} failed".format(i)
+
+
+def test_NamedTuple_init():
+    with pytest.raises(pt.TealInputError, match=r"NamedTuple must be subclassed$"):
+        abi.NamedTuple()
+
+    class Empty(abi.NamedTuple):
+        pass
+
+    with pytest.raises(
+        pt.TealInputError, match=r"Expected fields to be declared but found none$"
+    ):
+        Empty()
+
+    class ValidField(abi.NamedTuple):
+        name: abi.Field[abi.Uint16]
+
+    ValidField()
+
+    class NoField(abi.NamedTuple):
+        name: abi.Uint16
+
+    with pytest.raises(
+        pt.TealInputError,
+        match=r'Type annotation for attribute "name" must be a Field. Got ',
+    ):
+        NoField()
+
+
+class NT_0(abi.NamedTuple):
+    f0: abi.Field[abi.Uint64]
+    f1: abi.Field[abi.Uint32]
+    f2: abi.Field[abi.Uint16]
+    f3: abi.Field[abi.Uint8]
+
+
+class NT_1(abi.NamedTuple):
+    f0: abi.Field[abi.StaticArray[abi.Bool, Literal[4]]]
+    f1: abi.Field[abi.DynamicArray[abi.String]]
+    f2: abi.Field[abi.String]
+    f3: abi.Field[abi.Bool]
+    f4: abi.Field[abi.Address]
+    f5: abi.Field[NT_0]
+
+
+class NT_2(abi.NamedTuple):
+    f0: abi.Field[abi.Bool]
+    f1: abi.Field[abi.Bool]
+    f2: abi.Field[abi.Bool]
+    f3: abi.Field[abi.Bool]
+    f4: abi.Field[abi.Bool]
+    f5: abi.Field[abi.Bool]
+    f6: abi.Field[abi.Bool]
+    f7: abi.Field[abi.Bool]
+    f8: abi.Field[NT_1]
+
+
+class NT_3(abi.NamedTuple):
+    f0: abi.Field[NT_0]
+    f1: abi.Field[NT_1]
+    f2: abi.Field[NT_2]
+
+
+@pytest.mark.parametrize("test_case", [NT_0, NT_1, NT_2, NT_3])
+def test_NamedTuple_getitem(test_case: type[abi.NamedTuple]):
+    tuple_value = test_case()
+    tuple_len_static = tuple_value.type_spec().length_static()
+    for i in range(tuple_len_static):
+        elem_by_field: abi.TupleElement = getattr(tuple_value, f"f{i}")
+        elem_by_index: abi.TupleElement = tuple_value[i]
+
+        assert (
+            type(elem_by_field) is abi.TupleElement
+        ), f"Test case {test_case} at field f{i} must be TupleElement"
+        assert (
+            type(elem_by_index) is abi.TupleElement
+        ), f"Test case {test_case} at index {i} must be TupleElement"
+
+        assert (
+            elem_by_field.index == i
+        ), f"Test case {test_case} at field f{i} should have index {i}."
+        assert (
+            elem_by_index.index == i
+        ), f"Test case {test_case} at index {i} should have index {i}."
+
+        assert (
+            elem_by_field.tuple is tuple_value
+        ), f"Test case {test_case} at field f{i} should have attr tuple == {test_case}."
+        assert (
+            elem_by_index.tuple is tuple_value
+        ), f"Test case {test_case} at index {i} should have attr tuple == {test_case}."
+
+        assert (
+            elem_by_field.produced_type_spec() == elem_by_index.produced_type_spec()
+        ), f"Test case {test_case} at field f{i} type spec unmatching: {elem_by_field.produced_type_spec()} != {elem_by_index.produced_type_spec()}."
+
+    with pytest.raises(KeyError):
+        tuple_value.aaaaa
+
+    with pytest.raises(pt.TealInputError):
+        tuple_value.f0 = abi.Uint64()

pyteal/ast/abi/util.py

@@ -118,6 +118,8 @@ def type_spec_from_annotation(annotation: Any) -> TypeSpec:
         Tuple3,
         Tuple4,
         Tuple5,
+        NamedTuple,
+        NamedTupleTypeSpec,
     )
     from pyteal.ast.abi.string import StringTypeSpec, String
     from pyteal.ast.abi.address import AddressTypeSpec, Address
@@ -225,6 +227,9 @@ def type_spec_from_annotation(annotation: Any) -> TypeSpec:
     if origin is Tuple:
         return TupleTypeSpec(*(type_spec_from_annotation(arg) for arg in args))
 
+    if issubclass(origin, NamedTuple):
+        return cast(NamedTupleTypeSpec, origin().type_spec())
+
     if origin is Tuple0:
         if len(args) != 0:
             raise TypeError("Tuple0 expects 0 type arguments. Got: {}".format(args))