Source code for beaker.state.reserved

from abc import ABC, abstractmethod
from collections.abc import Callable
from typing import Generic, Literal, TypeAlias, TypeVar

from pyteal import Expr, SubroutineFnWrapper, TealType, abi
from pyteal.types import require_type

from beaker.consts import MAX_GLOBAL_STATE, MAX_LOCAL_STATE
from beaker.state._abc import (
    AppSpecSchemaFragment,
    GlobalStateStorage,
    LocalStateStorage,
    StateStorage,
)
from beaker.state.primitive import (
    GlobalStateValue,
    LocalStateValue,
    StateValue,
    identity_key_gen,
    prefix_key_gen,
)

__all__ = [
    "ReservedStateValue",
    "ReservedGlobalStateValue",
    "ReservedLocalStateValue",
]


KeyGenerator: TypeAlias = SubroutineFnWrapper | Callable[[Expr], Expr]


ST = TypeVar("ST", bound=StateValue)


class ReservedStateValue(Generic[ST], StateStorage, ABC):
    """Base Class for ReservedStateValues

    Attributes:
        stack_type (TealType): The type of the state value (either TealType.bytes or TealType.uint64)
        max_keys (int): Maximum number of keys to reserve for this reserved state value
        key_gen (subroutine): A subroutine returning TealType.bytes, used to create a key where some data is stored.
        descr (str): Description of the state value to provide some information to clients

    """

    def __init__(
        self,
        stack_type: Literal[TealType.bytes, TealType.uint64],
        max_keys: int,
        key_gen: KeyGenerator | None = None,
        descr: str | None = None,
        *,
        prefix: str | None = None,
    ):
        if stack_type not in (TealType.bytes, TealType.uint64):
            raise ValueError(f"Invalid stack type: {stack_type}")
        if prefix is not None:
            if key_gen is not None:
                raise ValueError("Only one of key_gen or prefix can be specified")
            if prefix:
                key_gen = prefix_key_gen(prefix)
            else:
                key_gen = identity_key_gen
        self.stack_type = stack_type
        self.max_keys = max_keys
        self.descr = descr
        self.key_gen = key_gen

    def __set_name__(self, owner: type, name: str) -> None:
        if self.key_gen is None:
            self.key_gen = prefix_key_gen(name)

    @property
    def key_gen(self) -> KeyGenerator | None:
        return self._key_gen

    @key_gen.setter
    def key_gen(self, value: KeyGenerator) -> None:
        if isinstance(value, SubroutineFnWrapper):
            require_type(value, TealType.bytes)
        self._key_gen = value

    def __getitem__(self, key_seed: Expr | abi.BaseType) -> ST:
        """Method to access the state value with the key seed provided"""
        key: Expr
        if isinstance(key_seed, abi.BaseType):
            key = key_seed.encode()
        else:
            key = key_seed

        if self.key_gen is not None:
            key = self.key_gen(key)
        return self._get_state_for_key(key)

    @abstractmethod
    def _get_state_for_key(self, key: Expr) -> ST:
        ...

    def num_keys(self) -> int:
        return self.max_keys

    def value_type(self) -> Literal[TealType.bytes, TealType.uint64]:
        return self.stack_type

    def app_spec_json(self) -> AppSpecSchemaFragment:
        return AppSpecSchemaFragment(
            "reserved",
            {
                "type": self.value_type().name,
                "max_keys": self.num_keys(),
                "descr": self.descr or "",
            },
        )


[docs]class ReservedGlobalStateValue( ReservedStateValue[GlobalStateValue], GlobalStateStorage ): """Reserved Application State (global state) Used when there should be a number of reserved state fields but the keys are uncertain at build time. Attributes: stack_type (TealType): The type of the state value (either TealType.bytes or TealType.uint64) max_keys (int): Maximum number of keys to reserve for this reserved state value key_gen (SubroutineFnWrapper): A subroutine returning TealType.bytes, used to create a key where some data is stored. descr (str): Description of the state value to provide some information to clients """
[docs] def __init__( self, stack_type: Literal[TealType.bytes, TealType.uint64], max_keys: int, key_gen: KeyGenerator | None = None, descr: str | None = None, *, prefix: str | None = None, ): super().__init__(stack_type, max_keys, key_gen, descr, prefix=prefix) if max_keys <= 0 or max_keys > MAX_GLOBAL_STATE: raise Exception(f"max keys expected to be between 0 and {MAX_GLOBAL_STATE}")
def initialize(self) -> Expr | None: return None def _get_state_for_key(self, key: Expr) -> GlobalStateValue: """Method to access the state value with the key seed provided""" return GlobalStateValue(stack_type=self.stack_type, key=key, descr=self.descr)
[docs]class ReservedLocalStateValue(ReservedStateValue[LocalStateValue], LocalStateStorage): """Reserved Account State (local state) Used when there should be a number of reserved state fields but the keys are uncertain at build time. Attributes: stack_type (TealType): The type of the state value (either TealType.bytes or TealType.uint64) max_keys (int): Maximum number of keys to reserve for this reserved state value key_gen (SubroutineFnWrapper): A subroutine returning TealType.bytes, used to create a key where some data is stored. descr (str): Description of the state value to provide some information to clients """
[docs] def __init__( self, stack_type: Literal[TealType.bytes, TealType.uint64], max_keys: int, key_gen: KeyGenerator | None = None, descr: str | None = None, *, prefix: str | None = None, ): super().__init__(stack_type, max_keys, key_gen, descr, prefix=prefix) if max_keys <= 0 or max_keys > MAX_LOCAL_STATE: raise Exception(f"max keys expected to be between 0 and {MAX_LOCAL_STATE}")
def initialize(self, acct: Expr) -> Expr | None: return None def _get_state_for_key(self, key: Expr) -> LocalStateValue: """Access AccountState value given key_seed""" return LocalStateValue(stack_type=self.stack_type, key=key, descr=self.descr)