Skip to content
Merged
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
24 changes: 20 additions & 4 deletions pyteal/ast/router.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ def __init__(self, name: Optional[str] = None) -> None:

@staticmethod
def __parse_conditions(
method_signature: str,
method_to_register: ABIReturnSubroutine | None,
on_completes: list[EnumInt],
creation: bool,
Expand All @@ -79,13 +80,18 @@ def __parse_conditions(
)
clear_state_conds: list[Expr] = []

if method_signature == "" and method_to_register is not None:
raise TealInputError(
"A method_signature must only be provided if method_to_register is not None"
)

# Check:
# - if current condition is for *ABI METHOD*
# (method selector && numAppArg == 1 + min(METHOD_APP_ARG_NUM_LIMIT, subroutineSyntaxArgNum))
# - or *BARE APP CALL* (numAppArg == 0)
method_or_bare_condition = (
And(
Txn.application_args[0] == MethodSignature(method_to_register.name()),
Txn.application_args[0] == MethodSignature(method_signature),
Txn.application_args.length()
== Int(
1
Expand Down Expand Up @@ -281,23 +287,33 @@ def on_bare_app_call(
else [cast(EnumInt, on_completes)]
)
approval_conds, clear_state_conds = Router.__parse_conditions(
method_to_register=None, on_completes=ocList, creation=creation
method_signature="",
method_to_register=None,
on_completes=ocList,
creation=creation,
)
branch = Router.__wrap_handler(False, bare_app_call)
self.__append_to_ast(approval_conds, clear_state_conds, branch, None)

def on_method_call(
self,
method_signature: str,
method_app_call: ABIReturnSubroutine,
*,
method_signature: str = None,
on_complete: EnumInt = OnComplete.NoOp,
creation: bool = False,
) -> None:
""" """
oc_list: list[EnumInt] = [cast(EnumInt, on_complete)]

if method_signature is None:
method_signature = method_app_call.method_signature()

approval_conds, clear_state_conds = Router.__parse_conditions(
method_to_register=method_app_call, on_completes=oc_list, creation=creation
method_signature=method_signature,
method_to_register=method_app_call,
on_completes=oc_list,
creation=creation,
)
branch = Router.__wrap_handler(True, method_app_call)
self.__append_to_ast(
Expand Down
9 changes: 9 additions & 0 deletions pyteal/ast/subroutine.py
Original file line number Diff line number Diff line change
Expand Up @@ -573,6 +573,15 @@ def __call__(
def name(self) -> str:
return self.subroutine.name()

def method_signature(self) -> str:
if not self.is_registrable():
raise TealInputError(
"Only registrable methods may return a method signature"
)

args = [str(v) for v in self.subroutine.abi_args.values()]
return f"{self.name()}({','.join(args)}){self.type_of()}"

def type_of(self) -> str | abi.TypeSpec:
return (
"void"
Expand Down
32 changes: 28 additions & 4 deletions pyteal/ast/subroutine_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ class ABISubroutineTC:
arg_instances: list[pt.Expr | pt.abi.BaseType]
name: str
ret_type: str | pt.abi.TypeSpec
signature: str


def test_abi_subroutine_definition():
Expand Down Expand Up @@ -130,13 +131,27 @@ def fn_2arg_1ret_with_expr(
return output.set(b[a % pt.Int(10)])

cases = (
ABISubroutineTC(fn_0arg_0ret, [], "fn_0arg_0ret", "void"),
ABISubroutineTC(fn_0arg_0ret, [], "fn_0arg_0ret", "void", "fn_0arg_0ret()void"),
ABISubroutineTC(
fn_0arg_uint64_ret, [], "fn_0arg_uint64_ret", pt.abi.Uint64TypeSpec()
fn_0arg_uint64_ret,
[],
"fn_0arg_uint64_ret",
pt.abi.Uint64TypeSpec(),
"fn_0arg_uint64_ret()uint64",
),
ABISubroutineTC(fn_1arg_0ret, [pt.abi.Uint64()], "fn_1arg_0ret", "void"),
ABISubroutineTC(
fn_1arg_1ret, [pt.abi.Uint64()], "fn_1arg_1ret", pt.abi.Uint64TypeSpec()
fn_1arg_0ret,
[pt.abi.Uint64()],
"fn_1arg_0ret",
"void",
"fn_1arg_0ret(uint64)void",
),
ABISubroutineTC(
fn_1arg_1ret,
[pt.abi.Uint64()],
"fn_1arg_1ret",
pt.abi.Uint64TypeSpec(),
"fn_1arg_1ret(uint64)uint64",
),
ABISubroutineTC(
fn_2arg_0ret,
Expand All @@ -148,6 +163,7 @@ def fn_2arg_1ret_with_expr(
],
"fn_2arg_0ret",
"void",
"fn_2arg_0ret(uint64,byte[10])void",
),
ABISubroutineTC(
fn_2arg_1ret,
Expand All @@ -159,6 +175,7 @@ def fn_2arg_1ret_with_expr(
],
"fn_2arg_1ret",
pt.abi.ByteTypeSpec(),
"fn_2arg_1ret(uint64,byte[10])byte",
),
ABISubroutineTC(
fn_2arg_1ret_with_expr,
Expand All @@ -170,6 +187,7 @@ def fn_2arg_1ret_with_expr(
],
"fn_2arg_1ret_with_expr",
pt.abi.ByteTypeSpec(),
None,
),
)

Expand All @@ -193,6 +211,12 @@ def fn_2arg_1ret_with_expr(
map(lambda x: isinstance(x, pt.abi.BaseType), case.arg_instances)
)

if case.definition.is_registrable():
assert case.definition.method_signature() == case.signature
else:
with pytest.raises(pt.TealInputError):
case.definition.method_signature()


def test_subroutine_definition_validate():
"""
Expand Down