method pseudo-op support for ABI methods

pyteal/ast/__init__.py

@@ -6,6 +6,7 @@
 from .addr import Addr
 from .bytes import Bytes
 from .int import Int, EnumInt
+from .method import Method
 
 # properties
 from .arg import Arg
@@ -126,6 +127,7 @@
     "Bytes",
     "Int",
     "EnumInt",
+    "Method",
     "Arg",
     "TxnType",
     "TxnField",

pyteal/ast/method.py

@@ -0,0 +1,43 @@
+from typing import TYPE_CHECKING
+from pyteal.errors import TealInputError
+
+from pyteal.types import TealType
+
+from ..types import TealType
+from ..ir import TealOp, Op, TealBlock
+from .leafexpr import LeafExpr
+
+if TYPE_CHECKING:
+    from ..compiler import CompileOptions
+
+
+class Method(LeafExpr):
+    """An expression that represents an ABI method selector"""
+
+    def __init__(self, methodName: str) -> None:
+        """Create a new method selector for ABI method call.
+
+        Args:
+            methodName: A string containing a valid ABI method signature
+        """
+        super().__init__()
+        if type(methodName) is not str:
+            raise TealInputError(
+                "invalid input type {} to Method".format(type(methodName))
+            )
+        elif len(methodName) == 0:
+            raise TealInputError("invalid input empty string to Method")
+        self.methodName = methodName
+
+    def __teal__(self, options: "CompileOptions"):
+        op = TealOp(self, Op.method, self.methodName)
+        return TealBlock.FromOp(options, op)
+
+    def __str__(self) -> str:
+        return "(method: {})".format(self.methodName)
+
+    def type_of(self) -> TealType:
+        return TealType.bytes
+
+
+Method.__module__ = "pyteal"

pyteal/ast/method_test.py

@@ -0,0 +1,25 @@
+import pytest
+
+from pyteal.ast.method import Method
+
+from .. import *
+
+
+def test_method():
+    expr = Method("add(uint64,uint64)uint64")
+    assert expr.type_of() == TealType.bytes
+
+    expected = TealSimpleBlock([TealOp(expr, Op.method, "add(uint64,uint64)uint64")])
+    actual, _ = expr.__teal__(CompileOptions())
+    assert expected == actual
+
+
+def test_method_invalid():
+    with pytest.raises(TealInputError):
+        Method(114514)
+
+    with pytest.raises(TealInputError):
+        Method(["m0()void", "m1()uint64"])
+
+    with pytest.raises(TealInputError):
+        Method("")

pyteal/ast/subroutine.py

@@ -10,7 +10,6 @@
 from .scratchvar import ScratchVar
 
 if TYPE_CHECKING:
-    from ..ir import TealSimpleBlock
     from ..compiler import CompileOptions
 
 
@@ -19,7 +18,10 @@ class SubroutineDefinition:
     nextSubroutineId = 0
 
     def __init__(
-        self, implementation: Callable[..., Expr], returnType: TealType
+        self,
+        implementation: Callable[..., Expr],
+        returnType: TealType,
+        methodSignature: str = None,
     ) -> None:
         super().__init__()
         self.id = SubroutineDefinition.nextSubroutineId
@@ -53,6 +55,7 @@ def __init__(
         self.returnType = returnType
 
         self.declaration: Optional["SubroutineDeclaration"] = None
+        self.methodSignature = methodSignature
 
     def getDeclaration(self) -> "SubroutineDeclaration":
         if self.declaration is None:
@@ -61,7 +64,11 @@ def getDeclaration(self) -> "SubroutineDeclaration":
         return self.declaration
 
     def name(self) -> str:
-        return self.implementation.__name__
+        return (
+            self.implementation.__name__
+            if self.methodSignature is None
+            else self.methodSignature
+        )
 
     def argumentCount(self) -> int:
         return len(self.implementationParams)
@@ -181,17 +188,20 @@ def mySubroutine(a: Expr, b: Expr) -> Expr:
             ])
     """
 
-    def __init__(self, returnType: TealType) -> None:
+    def __init__(self, returnType: TealType, methodSignature: str = None) -> None:
         """Define a new subroutine with the given return type.
 
         Args:
             returnType: The type that the return value of this subroutine must conform to.
                 TealType.none indicates that this subroutine does not return any value.
         """
         self.returnType = returnType
+        self.methodSignature = methodSignature
 
     def __call__(self, fnImplementation: Callable[..., Expr]) -> Callable[..., Expr]:
-        subroutine = SubroutineDefinition(fnImplementation, self.returnType)
+        subroutine = SubroutineDefinition(
+            fnImplementation, self.returnType, self.methodSignature
+        )
 
         @wraps(fnImplementation)
         def subroutineCall(*args: Expr, **kwargs) -> Expr:

pyteal/ir/ops.py

@@ -74,6 +74,7 @@ def min_version(self) -> int:
     bytec_3             = OpType("bytec_3",             Mode.Signature | Mode.Application, 2)
     byte                = OpType("byte",                Mode.Signature | Mode.Application, 2)
     addr                = OpType("addr",                Mode.Signature | Mode.Application, 2)
+    method              = OpType("method",              Mode.Signature | Mode.Application, 2)
     arg                 = OpType("arg",                 Mode.Signature,                    2)
     txn                 = OpType("txn",                 Mode.Signature | Mode.Application, 2)
     global_             = OpType("global",              Mode.Signature | Mode.Application, 2)

pyteal/ir/tealop.py

@@ -1,4 +1,4 @@
-from typing import cast, Union, List, Optional, TYPE_CHECKING
+from typing import Union, List, Optional, TYPE_CHECKING
 
 from .tealcomponent import TealComponent
 from .labelref import LabelReference