Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow various methods to accept a union of Expr, int, and/or str #152

Open
wants to merge 8 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions mypy.ini
Original file line number Diff line number Diff line change
Expand Up @@ -5,3 +5,6 @@ ignore_missing_imports = True

[mypy-algosdk.*]
ignore_missing_imports = True

[mypy-pyteal.*]
allow_redefinition = True
2 changes: 2 additions & 0 deletions pyteal/ast/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,7 @@
from .scratch import ScratchSlot, ScratchLoad, ScratchStore, ScratchStackStore
from .scratchvar import ScratchVar
from .maybe import MaybeValue
from .get_teal_type import get_teal_type

__all__ = [
"Expr",
Expand Down Expand Up @@ -240,4 +241,5 @@
"For",
"Break",
"Continue",
"get_teal_type",
]
87 changes: 68 additions & 19 deletions pyteal/ast/app.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Union
from enum import Enum

from ..types import TealType, require_type
Expand All @@ -8,6 +8,7 @@
from .maybe import MaybeValue
from .int import EnumInt
from .global_ import Global
from .get_teal_type import get_teal_type

if TYPE_CHECKING:
from ..compiler import CompileOptions
Expand Down Expand Up @@ -84,7 +85,9 @@ def id(cls) -> Global:
return Global.current_application_id()

@classmethod
def optedIn(cls, account: Expr, app: Expr) -> "App":
def optedIn(
cls, account: Union[int, str, Expr], app: Union[int, str, Expr]
) -> "App":
"""Check if an account has opted in for an application.

Args:
Expand All @@ -95,12 +98,14 @@ def optedIn(cls, account: Expr, app: Expr) -> "App":
must be evaluated to uint64 (or, since v4, an application id that appears in
Txn.ForeignApps or is the CurrentApplicationID, must be evaluated to bytes).
"""
account, app = map(get_teal_type, [account, app])

require_type(account, TealType.anytype)
require_type(app, TealType.uint64)
return cls(AppField.optedIn, [account, app])

@classmethod
def localGet(cls, account: Expr, key: Expr) -> "App":
def localGet(cls, account: Union[int, str, Expr], key: Union[str, Expr]) -> "App":
"""Read from an account's local state for the current application.

Args:
Expand All @@ -109,12 +114,19 @@ def localGet(cls, account: Expr, key: Expr) -> "App":
Txn.Accounts or is Txn.Sender, must be evaluated to bytes).
key: The key to read from the account's local state. Must evaluate to bytes.
"""
account, key = map(get_teal_type, [account, key])

require_type(account, TealType.anytype)
require_type(key, TealType.bytes)
return cls(AppField.localGet, [account, key])

@classmethod
def localGetEx(cls, account: Expr, app: Expr, key: Expr) -> MaybeValue:
def localGetEx(
cls,
account: Union[int, str, Expr],
app: Union[int, str, Expr],
key: Union[str, Expr],
) -> MaybeValue:
"""Read from an account's local state for an application.

Args:
Expand All @@ -126,6 +138,8 @@ def localGetEx(cls, account: Expr, app: Expr, key: Expr) -> MaybeValue:
Txn.ForeignApps or is the CurrentApplicationID, must be evaluated to bytes).
key: The key to read from the account's local state. Must evaluate to bytes.
"""
account, app, key = map(get_teal_type, [account, app, key])

require_type(account, TealType.anytype)
require_type(app, TealType.uint64)
require_type(key, TealType.bytes)
Expand All @@ -134,17 +148,19 @@ def localGetEx(cls, account: Expr, app: Expr, key: Expr) -> MaybeValue:
)

@classmethod
def globalGet(cls, key: Expr) -> "App":
def globalGet(cls, key: Union[str, Expr]) -> "App":
"""Read from the global state of the current application.

Args:
key: The key to read from the global application state. Must evaluate to bytes.
"""
key = get_teal_type(key)

require_type(key, TealType.bytes)
return cls(AppField.globalGet, [key])

@classmethod
def globalGetEx(cls, app: Expr, key: Expr) -> MaybeValue:
def globalGetEx(cls, app: Union[str, Expr], key: Union[str, Expr]) -> MaybeValue:
"""Read from the global state of an application.

Args:
Expand All @@ -153,14 +169,21 @@ def globalGetEx(cls, app: Expr, key: Expr) -> MaybeValue:
Txn.ForeignApps or is the CurrentApplicationID, must be evaluated to uint64).
key: The key to read from the global application state. Must evaluate to bytes.
"""
app, key = map(get_teal_type, [app, key])

require_type(app, TealType.uint64)
require_type(key, TealType.bytes)
return MaybeValue(
AppField.globalGetEx.get_op(), TealType.anytype, args=[app, key]
)

@classmethod
def localPut(cls, account: Expr, key: Expr, value: Expr) -> "App":
def localPut(
cls,
account: Union[int, str, Expr],
key: Union[str, Expr],
value: Union[int, str, Expr],
) -> "App":
"""Write to an account's local state for the current application.

Args:
Expand All @@ -170,25 +193,29 @@ def localPut(cls, account: Expr, key: Expr, value: Expr) -> "App":
key: The key to write in the account's local state. Must evaluate to bytes.
value: The value to write in the account's local state. Can evaluate to any type.
"""
account, value, key = map(get_teal_type, [account, value, key])

require_type(account, TealType.anytype)
require_type(key, TealType.bytes)
require_type(value, TealType.anytype)
return cls(AppField.localPut, [account, key, value])

@classmethod
def globalPut(cls, key: Expr, value: Expr) -> "App":
def globalPut(cls, key: Union[str, Expr], value: Union[int, str, Expr]) -> "App":
"""Write to the global state of the current application.

Args:
key: The key to write in the global application state. Must evaluate to bytes.
value: THe value to write in the global application state. Can evaluate to any type.
"""
value, key = map(get_teal_type, [value, key])

require_type(key, TealType.bytes)
require_type(value, TealType.anytype)
return cls(AppField.globalPut, [key, value])

@classmethod
def localDel(cls, account: Expr, key: Expr) -> "App":
def localDel(cls, account: Union[int, str, Expr], key: Union[str, Expr]) -> "App":
"""Delete a key from an account's local state for the current application.

Args:
Expand All @@ -197,17 +224,21 @@ def localDel(cls, account: Expr, key: Expr) -> "App":
Txn.Accounts or is Txn.Sender, must be evaluated to bytes).
key: The key to delete from the account's local state. Must evaluate to bytes.
"""
account, key = map(get_teal_type, [account, key])

require_type(account, TealType.anytype)
require_type(key, TealType.bytes)
return cls(AppField.localDel, [account, key])

@classmethod
def globalDel(cls, key: Expr) -> "App":
def globalDel(cls, key: Union[str, Expr]) -> "App":
"""Delete a key from the global state of the current application.

Args:
key: The key to delete from the global application state. Must evaluate to bytes.
"""
key = get_teal_type(key)

require_type(key, TealType.bytes)
return cls(AppField.globalDel, [key])

Expand All @@ -217,13 +248,15 @@ def globalDel(cls, key: Expr) -> "App":

class AppParam:
@classmethod
def approvalProgram(cls, app: Expr) -> MaybeValue:
def approvalProgram(cls, app: Union[int, Expr]) -> MaybeValue:
"""Get the bytecode of Approval Program for the application.

Args:
app: An index into Txn.ForeignApps that correspond to the application to check.
Must evaluate to uint64.
"""
app = get_teal_type(app)

require_type(app, TealType.uint64)
return MaybeValue(
Op.app_params_get,
Expand All @@ -233,13 +266,15 @@ def approvalProgram(cls, app: Expr) -> MaybeValue:
)

@classmethod
def clearStateProgram(cls, app: Expr) -> MaybeValue:
def clearStateProgram(cls, app: Union[int, Expr]) -> MaybeValue:
"""Get the bytecode of Clear State Program for the application.

Args:
app: An index into Txn.ForeignApps that correspond to the application to check.
Must evaluate to uint64.
"""
app = get_teal_type(app)

require_type(app, TealType.uint64)
return MaybeValue(
Op.app_params_get,
Expand All @@ -249,13 +284,15 @@ def clearStateProgram(cls, app: Expr) -> MaybeValue:
)

@classmethod
def globalNumUnit(cls, app: Expr) -> MaybeValue:
def globalNumUnit(cls, app: Union[int, Expr]) -> MaybeValue:
"""Get the number of uint64 values allowed in Global State for the application.

Args:
app: An index into Txn.ForeignApps that correspond to the application to check.
Must evaluate to uint64.
"""
app = get_teal_type(app)

require_type(app, TealType.uint64)
return MaybeValue(
Op.app_params_get,
Expand All @@ -265,13 +302,15 @@ def globalNumUnit(cls, app: Expr) -> MaybeValue:
)

@classmethod
def globalNumByteSlice(cls, app: Expr) -> MaybeValue:
def globalNumByteSlice(cls, app: Union[int, Expr]) -> MaybeValue:
"""Get the number of byte array values allowed in Global State for the application.

Args:
app: An index into Txn.ForeignApps that correspond to the application to check.
Must evaluate to uint64.
"""
app = get_teal_type(app)

require_type(app, TealType.uint64)
return MaybeValue(
Op.app_params_get,
Expand All @@ -281,13 +320,15 @@ def globalNumByteSlice(cls, app: Expr) -> MaybeValue:
)

@classmethod
def localNumUnit(cls, app: Expr) -> MaybeValue:
def localNumUnit(cls, app: Union[int, Expr]) -> MaybeValue:
"""Get the number of uint64 values allowed in Local State for the application.

Args:
app: An index into Txn.ForeignApps that correspond to the application to check.
Must evaluate to uint64.
"""
app = get_teal_type(app)

require_type(app, TealType.uint64)
return MaybeValue(
Op.app_params_get,
Expand All @@ -297,13 +338,15 @@ def localNumUnit(cls, app: Expr) -> MaybeValue:
)

@classmethod
def localNumByteSlice(cls, app: Expr) -> MaybeValue:
def localNumByteSlice(cls, app: Union[int, Expr]) -> MaybeValue:
"""Get the number of byte array values allowed in Local State for the application.

Args:
app: An index into Txn.ForeignApps that correspond to the application to check.
Must evaluate to uint64.
"""
app = get_teal_type(app)

require_type(app, TealType.uint64)
return MaybeValue(
Op.app_params_get,
Expand All @@ -313,13 +356,15 @@ def localNumByteSlice(cls, app: Expr) -> MaybeValue:
)

@classmethod
def extraProgramPages(cls, app: Expr) -> MaybeValue:
def extraProgramPages(cls, app: Union[int, Expr]) -> MaybeValue:
"""Get the number of Extra Program Pages of code space for the application.

Args:
app: An index into Txn.ForeignApps that correspond to the application to check.
Must evaluate to uint64.
"""
app = get_teal_type(app)

require_type(app, TealType.uint64)
return MaybeValue(
Op.app_params_get,
Expand All @@ -329,26 +374,30 @@ def extraProgramPages(cls, app: Expr) -> MaybeValue:
)

@classmethod
def creator(cls, app: Expr) -> MaybeValue:
def creator(cls, app: Union[int, Expr]) -> MaybeValue:
"""Get the creator address for the application.

Args:
app: An index into Txn.ForeignApps that correspond to the application to check.
Must evaluate to uint64.
"""
app = get_teal_type(app)

require_type(app, TealType.uint64)
return MaybeValue(
Op.app_params_get, TealType.bytes, immediate_args=["AppCreator"], args=[app]
)

@classmethod
def address(cls, app: Expr) -> MaybeValue:
def address(cls, app: Union[int, Expr]) -> MaybeValue:
"""Get the escrow address for the application.

Args:
app: An index into Txn.ForeignApps that correspond to the application to check.
Must evaluate to uint64.
"""
app = get_teal_type(app)

require_type(app, TealType.uint64)
return MaybeValue(
Op.app_params_get, TealType.bytes, immediate_args=["AppAddress"], args=[app]
Expand Down
Loading