diff --git a/pyteal/ast/router.py b/pyteal/ast/router.py index 71b8f35b3..6297fb37e 100644 --- a/pyteal/ast/router.py +++ b/pyteal/ast/router.py @@ -6,7 +6,10 @@ from algosdk import encoding from pyteal.config import METHOD_ARG_NUM_CUTOFF -from pyteal.errors import TealInputError, TealInternalError +from pyteal.errors import ( + TealInputError, + TealInternalError, +) from pyteal.types import TealType from pyteal.compiler.compiler import compileTeal, DEFAULT_TEAL_VERSION, OptimizeOptions from pyteal.ir.ops import Mode @@ -447,6 +450,7 @@ def __init__( self.approval_ast = ASTBuilder() self.clear_state_ast = ASTBuilder() + self.methods: list[sdk_abi.Method] = [] self.method_sig_to_selector: dict[str, bytes] = dict() self.method_selector_to_sig: dict[bytes, str] = dict() @@ -473,6 +477,7 @@ def add_method_handler( method_call: ABIReturnSubroutine, overriding_name: str = None, method_config: MethodConfig = None, + description: str = None, ) -> None: if not isinstance(method_call, ABIReturnSubroutine): raise TealInputError( @@ -494,6 +499,12 @@ def add_method_handler( f"re-registering method {method_signature} has hash collision " f"with {self.method_selector_to_sig[method_selector]}" ) + + meth = method_call.method_spec() + if description is not None: + meth.desc = description + self.methods.append(meth) + self.method_sig_to_selector[method_signature] = method_selector self.method_selector_to_sig[method_selector] = method_signature @@ -518,6 +529,7 @@ def method( clear_state: CallConfig = None, update_application: CallConfig = None, delete_application: CallConfig = None, + description: str = None, ): """ A decorator style method registration by decorating over a python function, @@ -567,7 +579,7 @@ def none_to_never(x: None | CallConfig): update_application=_update_app, delete_application=_delete_app, ) - self.add_method_handler(wrapped_subroutine, name, call_configs) + self.add_method_handler(wrapped_subroutine, name, call_configs, description) if not func: return wrap @@ -576,19 +588,15 @@ def none_to_never(x: None | CallConfig): def contract_construct(self) -> sdk_abi.Contract: """A helper function in constructing contract JSON object. - It takes out the method signatures from approval program `ProgramNode`'s, + It takes out the method spec from approval program methods, and constructs an `Contract` object. Returns: contract: a dictified `Contract` object constructed from - approval program's method signatures and `self.name`. + approval program's method specs and `self.name`. """ - method_collections = [ - sdk_abi.Method.from_signature(sig) - for sig in self.method_sig_to_selector - if isinstance(sig, str) - ] - return sdk_abi.Contract(self.name, method_collections) + + return sdk_abi.Contract(self.name, self.methods) def build_program(self) -> tuple[Expr, Expr, sdk_abi.Contract]: """ diff --git a/pyteal/ast/router_test.py b/pyteal/ast/router_test.py index d9b6a967a..bba67c88b 100644 --- a/pyteal/ast/router_test.py +++ b/pyteal/ast/router_test.py @@ -10,11 +10,13 @@ @pt.ABIReturnSubroutine def add(a: pt.abi.Uint64, b: pt.abi.Uint64, *, output: pt.abi.Uint64) -> pt.Expr: + """add takes 2 integers a,b and adds them, returning the sum""" return output.set(a.get() + b.get()) @pt.ABIReturnSubroutine def sub(a: pt.abi.Uint64, b: pt.abi.Uint64, *, output: pt.abi.Uint64) -> pt.Expr: + """replace me""" return output.set(a.get() - b.get()) @@ -528,8 +530,29 @@ def test_contract_json_obj(): router = pt.Router(contract_name, on_complete_actions) method_list: list[sdk_abi.Method] = [] for subroutine in abi_subroutines: - router.add_method_handler(subroutine) - method_list.append(sdk_abi.Method.from_signature(subroutine.method_signature())) + + doc = subroutine.subroutine.implementation.__doc__ + desc = None + if doc is not None and doc.strip() == "replace me": + desc = "dope description" + + router.add_method_handler(subroutine, description=desc) + + ms = subroutine.method_spec() + + # Manually replace it since the override is applied in the method handler + # not attached to the ABIReturnSubroutine itself + ms.desc = desc if desc is not None else ms.desc + + sig_method = sdk_abi.Method.from_signature(subroutine.method_signature()) + + assert ms.name == sig_method.name + + for idx, arg in enumerate(ms.args): + assert arg.type == sig_method.args[idx].type + + method_list.append(ms) + sdk_contract = sdk_abi.Contract(contract_name, method_list) contract = router.contract_construct() assert contract == sdk_contract diff --git a/pyteal/ast/subroutine.py b/pyteal/ast/subroutine.py index 4a8dec550..0fbd7a217 100644 --- a/pyteal/ast/subroutine.py +++ b/pyteal/ast/subroutine.py @@ -3,6 +3,8 @@ from types import MappingProxyType, NoneType from typing import Any, Callable, Final, Optional, TYPE_CHECKING, cast +import algosdk.abi as sdk_abi + from pyteal.ast import abi from pyteal.ast.expr import Expr from pyteal.ast.seq import Seq @@ -611,6 +613,35 @@ def method_signature(self, overriding_name: str = None) -> str: overriding_name = self.name() return f"{overriding_name}({','.join(args)}){self.type_of()}" + def method_spec(self) -> sdk_abi.Method: + skip_names = ["return", "output"] + + args = [ + { + "type": str(abi.type_spec_from_annotation(val)), + "name": name, + } + for name, val in self.subroutine.annotations.items() + if name not in skip_names + ] + + spec = { + "name": self.name(), + "args": args, + "returns": {"type": str(self.type_of())}, + } + + if self.subroutine.implementation.__doc__ is not None: + spec["desc"] = " ".join( + [ + i.strip() + for i in self.subroutine.implementation.__doc__.split("\n") + if not (i.isspace() or len(i) == 0) + ] + ) + + return sdk_abi.Method.undictify(spec) + def type_of(self) -> str | abi.TypeSpec: return ( "void"