Method sig for router (#340)

* adding method_signature to ABIReturnSubroutine

Author: barnjamin

pyteal/ast/router.py

@@ -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,
@@ -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
@@ -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(

pyteal/ast/subroutine.py

@@ -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"

pyteal/ast/subroutine_test.py

@@ -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():
@@ -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,
@@ -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,
@@ -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,
@@ -170,6 +187,7 @@ def fn_2arg_1ret_with_expr(
             ],
             "fn_2arg_1ret_with_expr",
             pt.abi.ByteTypeSpec(),
+            None,
         ),
     )
 
@@ -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():
     """