From aae4548a59f135c5955fe90df11ebd6ca0e9a3b5 Mon Sep 17 00:00:00 2001
From: Zeph Grunschlag <tzaffi@users.noreply.github.com>
Date: Wed, 2 Feb 2022 12:52:30 -0600
Subject: [PATCH] Subroutine Type Annotations (#182)

This PR requires that any type annotation of a Subroutine parameter or return value be of type `Expr`. Missing annotations are assumed to be `Expr`'s as well. In a follow up PR #183 this restriction will be loosened.
---
 pyteal/ast/subroutine.py      | 10 +++++++
 pyteal/ast/subroutine_test.py | 55 ++++++++++++++++++++++++++++++-----
 2 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/pyteal/ast/subroutine.py b/pyteal/ast/subroutine.py
index 136be53c9..0bf688d28 100644
--- a/pyteal/ast/subroutine.py
+++ b/pyteal/ast/subroutine.py
@@ -49,6 +49,16 @@ def __init__(
                     )
                 )
 
+        for var, var_type in implementation.__annotations__.items():
+            if var_type is not Expr:
+                stub = "Return" if var == "return" else ("parameter " + var)
+
+                raise TealInputError(
+                    "Function has {} of disallowed type {}. Only type Expr is allowed".format(
+                        stub, var_type
+                    )
+                )
+
         self.implementation = implementation
         self.implementationParams = sig.parameters
         self.returnType = returnType
diff --git a/pyteal/ast/subroutine_test.py b/pyteal/ast/subroutine_test.py
index a283b3a55..c60da0907 100644
--- a/pyteal/ast/subroutine_test.py
+++ b/pyteal/ast/subroutine_test.py
@@ -28,6 +28,18 @@ def fn10Args(a1, a2, a3, a4, a5, a6, a7, a8, a9, a10):
     lam2Args = lambda a1, a2: Return()
     lam10Args = lambda a1, a2, a3, a4, a5, a6, a7, a8, a9, a10: Return()
 
+    def fnWithExprAnnotations(a: Expr, b: Expr) -> Expr:
+        return Return()
+
+    def fnWithOnlyReturnExprAnnotations(a, b) -> Expr:
+        return Return()
+
+    def fnWithOnlyArgExprAnnotations(a: Expr, b: Expr):
+        return Return()
+
+    def fnWithPartialExprAnnotations(a, b: Expr) -> Expr:
+        return Return()
+
     cases = (
         (fn0Args, 0, "fn0Args"),
         (fn1Args, 1, "fn1Args"),
@@ -37,6 +49,10 @@ def fn10Args(a1, a2, a3, a4, a5, a6, a7, a8, a9, a10):
         (lam1Args, 1, "<lambda>"),
         (lam2Args, 2, "<lambda>"),
         (lam10Args, 10, "<lambda>"),
+        (fnWithExprAnnotations, 2, "fnWithExprAnnotations"),
+        (fnWithOnlyReturnExprAnnotations, 2, "fnWithOnlyReturnExprAnnotations"),
+        (fnWithOnlyArgExprAnnotations, 2, "fnWithOnlyArgExprAnnotations"),
+        (fnWithPartialExprAnnotations, 2, "fnWithPartialExprAnnotations"),
     )
 
     for (fn, numArgs, name) in cases:
@@ -72,18 +88,43 @@ def fnWithKeywordArgs(a, *, b):
     def fnWithVariableArgs(a, *b):
         return Return()
 
+    def fnWithNonExprReturnAnnotation(a, b) -> TealType.uint64:
+        return Return()
+
+    def fnWithNonExprParamAnnotation(a, b: TealType.uint64):
+        return Return()
+
     cases = (
-        1,
-        None,
-        fnWithDefaults,
-        fnWithKeywordArgs,
-        fnWithVariableArgs,
+        (1, "TealInputError('Input to SubroutineDefinition is not callable'"),
+        (None, "TealInputError('Input to SubroutineDefinition is not callable'"),
+        (
+            fnWithDefaults,
+            "TealInputError('Function has a parameter with a default value, which is not allowed in a subroutine: b'",
+        ),
+        (
+            fnWithKeywordArgs,
+            "TealInputError('Function has a parameter type that is not allowed in a subroutine: parameter b with type KEYWORD_ONLY'",
+        ),
+        (
+            fnWithVariableArgs,
+            "TealInputError('Function has a parameter type that is not allowed in a subroutine: parameter b with type VAR_POSITIONAL'",
+        ),
+        (
+            fnWithNonExprReturnAnnotation,
+            "TealInputError('Function has Return of disallowed type TealType.uint64. Only type Expr is allowed'",
+        ),
+        (
+            fnWithNonExprParamAnnotation,
+            "TealInputError('Function has parameter b of disallowed type TealType.uint64. Only type Expr is allowed'",
+        ),
     )
 
-    for case in cases:
-        with pytest.raises(TealInputError):
+    for case, msg in cases:
+        with pytest.raises(TealInputError) as e:
             SubroutineDefinition(case, TealType.none)
 
+            assert msg in str(e), "failed for case [{}]".format(case)
+
 
 def test_subroutine_declaration():
     cases = (