define reference type spec

pyteal/__init__.pyi

@@ -30,18 +30,22 @@ from pyteal.config import (
 __all__ = [
     "ABIReturnSubroutine",
     "AccountParam",
+    "AccountParamObject",
     "Add",
     "Addr",
     "And",
     "App",
     "AppField",
     "AppParam",
+    "AppParamObject",
     "Approve",
     "Arg",
     "Array",
     "Assert",
     "AssetHolding",
+    "AssetHoldingObject",
     "AssetParam",
+    "AssetParamObject",
     "Balance",
     "BareCallActions",
     "BinaryExpr",

pyteal/ast/__init__.py

@@ -24,9 +24,14 @@
 from pyteal.ast.gitxn import Gitxn, GitxnExpr, GitxnaExpr, InnerTxnGroup
 from pyteal.ast.gload import ImportScratchValue
 from pyteal.ast.global_ import Global, GlobalField
-from pyteal.ast.app import App, AppField, OnComplete, AppParam
-from pyteal.ast.asset import AssetHolding, AssetParam
-from pyteal.ast.acct import AccountParam
+from pyteal.ast.app import App, AppField, OnComplete, AppParam, AppParamObject
+from pyteal.ast.asset import (
+    AssetHolding,
+    AssetHoldingObject,
+    AssetParam,
+    AssetParamObject,
+)
+from pyteal.ast.acct import AccountParam, AccountParamObject
 
 # inner txns
 from pyteal.ast.itxn import InnerTxnBuilder, InnerTxn, InnerTxnAction
@@ -177,9 +182,13 @@
     "AppField",
     "OnComplete",
     "AppParam",
+    "AppParamObject",
     "AssetHolding",
+    "AssetHoldingObject",
     "AssetParam",
+    "AssetParamObject",
     "AccountParam",
+    "AccountParamObject",
     "InnerTxnBuilder",
     "InnerTxn",
     "InnerTxnAction",

pyteal/ast/abi/reference_type.py

@@ -1,14 +1,84 @@
-from typing import List, Final
-from pyteal.ast.abi.type import TypeSpec
-from pyteal.ast.abi.uint import Uint, UintTypeSpec
+from typing import List, 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
+
 from pyteal.ast.expr import Expr
 from pyteal.ast.txn import Txn
+from pyteal.ast.acct import AccountParamObject
+from pyteal.ast.asset import AssetHoldingObject, AssetParamObject
+from pyteal.ast.app import AppParamObject
+from pyteal.errors import TealInputError
+from pyteal.types import TealType
+
+
+T = TypeVar("T", bound="ReferenceType")
+
+
+class ReferenceTypeSpec(TypeSpec):
+    @abstractmethod
+    def new_instance(self) -> "ReferenceType":
+        pass
+
+    @abstractmethod
+    def annotation_type(self) -> "type[ReferenceType]":
+        pass
+
+    def bit_size(self) -> int:
+        """Get the bit size of the index this reference type holds"""
+        return NUM_BITS_IN_BYTE
+
+    def is_dynamic(self) -> bool:
+        return False
+
+    def byte_length_static(self) -> int:
+        return 1
+
+    def storage_type(self) -> TealType:
+        return TealType.uint64
+
+
+class ReferenceType(BaseType):
+    @abstractmethod
+    def __init__(self, spec: ReferenceTypeSpec) -> None:
+        super().__init__(spec)
+
+    def type_spec(self) -> ReferenceTypeSpec:
+        return cast(ReferenceTypeSpec, super().type_spec())
 
+    def referenced_index(self) -> Expr:
+        """Get the reference index for this value.
 
-class AccountTypeSpec(UintTypeSpec):
-    def __init__(self):
-        super().__init__(8)
+        The three reference types (account, application, asset) contain indexes into a foreign array
+        of the transaction. This method returns that index.
 
+        If this reference type is an application or asset, note that this DOES NOT return the
+        application or asset ID. See :code:`application_id()` or :code:`asset_id()` for that.
+        """
+        return self.stored_value.load()
+
+    def decode(
+        self,
+        encoded: Expr,
+        *,
+        startIndex: Expr = None,
+        endIndex: Expr = None,
+        length: Expr = None,
+    ) -> Expr:
+        return uint_decode(
+            self.type_spec().bit_size(),
+            self.stored_value,
+            encoded,
+            startIndex,
+            endIndex,
+            length,
+        )
+
+    def encode(self) -> Expr:
+        raise TealInputError("A ReferenceType cannot be encoded")
+
+
+class AccountTypeSpec(ReferenceTypeSpec):
     def new_instance(self) -> "Account":
         return Account()
 
@@ -25,21 +95,39 @@ def __eq__(self, other: object) -> bool:
 AccountTypeSpec.__module__ = "pyteal"
 
 
-class Account(Uint):
+class Account(ReferenceType):
     def __init__(self) -> None:
         super().__init__(AccountTypeSpec())
 
-    def deref(self) -> Expr:
+    def address(self) -> Expr:
+        """Get the address of the account."""
         return Txn.accounts[self.stored_value.load()]
 
+    def params(self) -> AccountParamObject:
+        """Get information about the account."""
+        return AccountParamObject(self.referenced_index())
 
-Account.__module__ = "pyteal"
+    def asset_holding(self, asset: "Expr | Asset") -> AssetHoldingObject:
+        """Get information about an asset held by this account.
+
+        Args:
+            asset: An identifier for the asset. It must be one of the following: an abi.Asset
+                reference object, an expression holding an index into Txn.ForeignAssets that
+                corresponds to the asset (in which case it must evaluate to uint64), or since v4, an
+                expression holding an asset ID that appears in Txn.ForeignAssets (in which case it
+                must evaluate to uint64).
+        """
+        if isinstance(asset, Asset):
+            asset_ref = asset.referenced_index()
+        else:
+            asset_ref = asset
+        return AssetHoldingObject(asset_ref, self.referenced_index())
 
 
-class AssetTypeSpec(UintTypeSpec):
-    def __init__(self):
-        super().__init__(8)
+Account.__module__ = "pyteal"
 
+
+class AssetTypeSpec(ReferenceTypeSpec):
     def new_instance(self) -> "Asset":
         return Asset()
 
@@ -56,21 +144,39 @@ def __eq__(self, other: object) -> bool:
 AssetTypeSpec.__module__ = "pyteal"
 
 
-class Asset(Uint):
+class Asset(ReferenceType):
     def __init__(self) -> None:
         super().__init__(AssetTypeSpec())
 
-    def deref(self) -> Expr:
-        return Txn.assets[self.stored_value.load()]
+    def asset_id(self) -> Expr:
+        """Get the ID of the asset."""
+        return Txn.assets[self.referenced_index()]
 
+    def holding(self, account: Expr | Account) -> AssetHoldingObject:
+        """Get information about this asset held by an account.
 
-Asset.__module__ = "pyteal"
+        Args:
+            account: An identifier for the account. It must be one of the following: an abi.Account
+                reference object, an expression holding an index into Txn.Accounts that corresponds
+                to the account (in which case it must evaluate to uint64), or since v4, an
+                expression holding an account address that appears in Txn.Accounts or is Txn.Sender
+                (in which case it must evaluate to bytes).
+        """
+        if isinstance(account, Account):
+            account_ref = account.referenced_index()
+        else:
+            account_ref = account
+        return AssetHoldingObject(self.referenced_index(), account_ref)
+
+    def params(self) -> AssetParamObject:
+        """Get information about the asset's parameters."""
+        return AssetParamObject(self.referenced_index())
 
 
-class ApplicationTypeSpec(UintTypeSpec):
-    def __init__(self):
-        super().__init__(8)
+Asset.__module__ = "pyteal"
+
 
+class ApplicationTypeSpec(ReferenceTypeSpec):
     def new_instance(self) -> "Application":
         return Application()
 
@@ -87,13 +193,18 @@ def __eq__(self, other: object) -> bool:
 ApplicationTypeSpec.__module__ = "pyteal"
 
 
-class Application(Uint):
+class Application(ReferenceType):
     def __init__(self) -> None:
         super().__init__(ApplicationTypeSpec())
 
-    def deref(self) -> Expr:
+    def application_id(self) -> Expr:
+        """Get the ID of the application."""
         return Txn.applications[self.stored_value.load()]
 
+    def params(self) -> AppParamObject:
+        """Get information about the application's parameters."""
+        return AppParamObject(self.referenced_index())
+
 
 Application.__module__ = "pyteal"