ABI Method Return (#175)

Encode ABI variable and concat with `RETURN_METHOD_SELECTOR` and log it.

Author: ahangsu

pyteal/ast/abi/__init__.py

@@ -29,6 +29,7 @@
 from .array_static import StaticArrayTypeSpec, StaticArray
 from .array_dynamic import DynamicArrayTypeSpec, DynamicArray
 from .util import type_spec_from_annotation
+from .method_return import MethodReturn
 
 __all__ = [
     "TypeSpec",
@@ -65,4 +66,5 @@
     "DynamicArrayTypeSpec",
     "DynamicArray",
     "type_spec_from_annotation",
+    "MethodReturn",
 ]

pyteal/ast/abi/method_return.py

@@ -0,0 +1,39 @@
+from typing import TYPE_CHECKING, Tuple
+from . import BaseType
+from ...types import TealType
+from ...errors import TealInputError
+from .. import Expr, Log, Concat, Bytes
+from ...ir import TealBlock, TealSimpleBlock, Op
+from ...config import RETURN_METHOD_SELECTOR
+
+if TYPE_CHECKING:
+    from ...compiler import CompileOptions
+
+
+class MethodReturn(Expr):
+    def __init__(self, arg: BaseType):
+        super().__init__()
+        if not isinstance(arg, BaseType):
+            raise TealInputError(f"Expecting an ABI type argument but get {arg}")
+        self.arg = arg
+
+    def __teal__(self, options: "CompileOptions") -> Tuple[TealBlock, TealSimpleBlock]:
+        if options.version < Op.log.min_version:
+            raise TealInputError(
+                f"current version {options.version} is lower than log's min version {Op.log.min_version}"
+            )
+        return Log(
+            Concat(Bytes("base16", RETURN_METHOD_SELECTOR), self.arg.encode())
+        ).__teal__(options)
+
+    def __str__(self) -> str:
+        return f"(MethodReturn {self.arg.type_spec()})"
+
+    def type_of(self) -> TealType:
+        return TealType.none
+
+    def has_return(self) -> bool:
+        return False
+
+
+MethodReturn.__module__ = "pyteal"

pyteal/ast/abi/method_return_test.py

@@ -0,0 +1,35 @@
+import pytest
+
+from . import *
+from .. import Int, Bytes
+from ...types import TealType
+from ...errors import TealInputError
+
+
+POSITIVE_CASES = [
+    Uint16(),
+    Uint32(),
+    StaticArray(BoolTypeSpec(), 12),
+]
+
+
+@pytest.mark.parametrize("case", POSITIVE_CASES)
+def test_method_return(case):
+    m_ret = MethodReturn(case)
+    assert m_ret.type_of() == TealType.none
+    assert not m_ret.has_return()
+    assert str(m_ret) == f"(MethodReturn {case.type_spec()})"
+
+
+NEGATIVE_CASES = [
+    Int(0),
+    Bytes("aaaaaaa"),
+    Uint16,
+    Uint32,
+]
+
+
+@pytest.mark.parametrize("case", NEGATIVE_CASES)
+def test_method_return_error(case):
+    with pytest.raises(TealInputError):
+        MethodReturn(case)

pyteal/config.py

@@ -3,3 +3,6 @@
 
 # Number of scratch space slots available.
 NUM_SLOTS = 256
+
+# Method return selector in base16
+RETURN_METHOD_SELECTOR = "0x151F7C75"