Merge pull request #1 from michaeldiamant/more-teal6-ops_gitxnas_itxnas_validation

michaeldiamant · web-flow · commit fdb483368d34 · 2022-02-16T17:09:06.000-05:00
Move type validation to constructors for gtxn and gitxn variants
diff --git a/pyteal/ast/gitxn.py b/pyteal/ast/gitxn.py
@@ -6,6 +6,7 @@
 from ..ir import TealOp, Op, TealBlock
 from .expr import Expr
 from .txn import TxnExpr, TxnField, TxnObject, TxnaExpr
+from ..types import require_type, TealType
 
 if TYPE_CHECKING:
     from ..compiler import CompileOptions
@@ -16,6 +17,15 @@ class GitxnExpr(TxnExpr):
 
     def __init__(self, txnIndex: int, field: TxnField) -> None:
         super().__init__(Op.gitxn, "Gitxn", field)
+
+        # Currently we do not have gitxns. Only gitxn with immediate transaction index supported.
+        if type(txnIndex) is not int:
+            raise TealInputError(
+                "Invalid gitxn syntax with immediate transaction field {} and transaction index {}".format(
+                    field, txnIndex
+                )
+            )
+
         self.txnIndex = txnIndex
 
     def __str__(self):
@@ -24,14 +34,6 @@ def __str__(self):
     def __teal__(self, options: "CompileOptions"):
         verifyFieldVersion(self.field.arg_name, self.field.min_version, options.version)
 
-        # currently we do not have gitxns, only gitxn with immediate transaction index supported
-        if type(self.txnIndex) is not int:
-            raise TealInputError(
-                "Invalid gitxn syntax with immediate transaction field {} and transaction index {}".format(
-                    self.field, self.txnIndex
-                )
-            )
-
         verifyTealVersion(
             Op.gitxn.min_version,
             options.version,
@@ -49,6 +51,12 @@ class GitxnaExpr(TxnaExpr):
 
     def __init__(self, txnIndex: int, field: TxnField, index: Union[int, Expr]) -> None:
         super().__init__(Op.gitxna, Op.gitxnas, "Gitxna", field, index)
+
+        if type(txnIndex) is not int:
+            raise TealInputError(
+                f"Invalid txnIndex type:  Expected int, but received {txnIndex}."
+            )
+
         self.txnIndex = txnIndex
 
     def __str__(self):
@@ -59,15 +67,6 @@ def __str__(self):
     def __teal__(self, options: "CompileOptions"):
         verifyFieldVersion(self.field.arg_name, self.field.min_version, options.version)
 
-        if type(self.txnIndex) is not int or not (
-            type(self.index) is int or isinstance(self.index, Expr)
-        ):
-            raise TealInputError(
-                "Invalid gitxna syntax with immediate transaction index {}, transaction field {}, array index {}".format(
-                    self.txnIndex, self.field, self.index
-                )
-            )
-
         if type(self.index) is int:
             opToUse = Op.gitxna
         else:
diff --git a/pyteal/ast/gitxn_test.py b/pyteal/ast/gitxn_test.py
@@ -7,35 +7,75 @@
 
 
 def test_gitxn_invalid():
-    with pytest.raises(TealInputError):
-        GitxnExpr(0, TxnField.sender).__teal__(teal5Options)
-
-    with pytest.raises(TealInputError):
-        Gitxn[MAX_GROUP_SIZE].sender()
-
-    with pytest.raises(TealInputError):
-        Gitxn[-1].asset_sender()
-
-    with pytest.raises(TealInputError):
-        Gitxn[Bytes("first")].sender()
+    for ctor, e in [
+        (
+            lambda: Gitxn[MAX_GROUP_SIZE],
+            TealInputError,
+        ),
+        (
+            lambda: Gitxn[-1],
+            TealInputError,
+        ),
+    ]:
+        with pytest.raises(e):
+            ctor()
 
 
 def test_gitxn_valid():
     for i in range(MAX_GROUP_SIZE):
         Gitxn[i].sender()
 
 
-def test_gitxna_invalid():
-    with pytest.raises(TealInputError):
-        GitxnaExpr("Invalid_type", TxnField.application_args, 1).__teal__(teal6Options)
-
-    with pytest.raises(TealInputError):
-        GitxnaExpr(0, TxnField.application_args, "Invalid_type").__teal__(teal6Options)
+def test_gitxn_expr_invalid():
+    for f, e in [
+        (
+            lambda: GitxnExpr(Int(1), TxnField.sender),
+            TealInputError,
+        ),
+        (
+            lambda: GitxnExpr(1, TxnField.sender).__teal__(teal5Options),
+            TealInputError,
+        ),
+    ]:
+        with pytest.raises(e):
+            f()
+
+
+def test_gitxn_expr_valid():
+    GitxnExpr(1, TxnField.sender).__teal__(teal6Options)
+
+
+def test_gitxna_expr_invalid():
+    for f, e in [
+        (
+            lambda: GitxnaExpr("Invalid_type", TxnField.application_args, 1),
+            TealInputError,
+        ),
+        (
+            lambda: GitxnaExpr(1, TxnField.application_args, "Invalid_type"),
+            TealInputError,
+        ),
+        (
+            lambda: GitxnaExpr(0, TxnField.application_args, Assert(Int(1))),
+            TealTypeError,
+        ),
+        (
+            lambda: GitxnaExpr(0, TxnField.application_args, 0).__teal__(teal5Options),
+            TealInputError,
+        ),
+    ]:
+        with pytest.raises(e):
+            f()
 
 
 def test_gitxna_valid():
-    GitxnaExpr(0, TxnField.application_args, 1).__teal__(teal6Options)
-    GitxnaExpr(0, TxnField.application_args, Assert(Int(1))).__teal__(teal6Options)
+    [
+        e.__teal__(teal6Options)
+        for e in [
+            GitxnaExpr(0, TxnField.application_args, 1),
+            GitxnaExpr(0, TxnField.application_args, Int(1)),
+        ]
+    ]
 
 
 # txn_test.py performs additional testing
diff --git a/pyteal/ast/gtxn.py b/pyteal/ast/gtxn.py
@@ -11,11 +11,21 @@
     from ..compiler import CompileOptions
 
 
+def validate_txn_index_or_throw(txnIndex: Union[int, Expr]):
+    if not isinstance(txnIndex, (int, Expr)):
+        raise TealInputError(
+            f"Invalid txnIndex type:  Expected int or Expr, but received {txnIndex}"
+        )
+    if isinstance(txnIndex, Expr):
+        require_type(txnIndex, TealType.uint64)
+
+
 class GtxnExpr(TxnExpr):
     """An expression that accesses a transaction field from a transaction in the current group."""
 
     def __init__(self, txnIndex: Union[int, Expr], field: TxnField) -> None:
         super().__init__(Op.gtxn, "Gtxn", field)
+        validate_txn_index_or_throw(txnIndex)
         self.txnIndex = txnIndex
 
     def __str__(self):
@@ -51,6 +61,7 @@ def __init__(
         self, txnIndex: Union[int, Expr], field: TxnField, index: Union[int, Expr]
     ) -> None:
         super().__init__(Op.gtxna, Op.gtxnas, "Gtxna", field, index)
+        validate_txn_index_or_throw(txnIndex)
         self.txnIndex = txnIndex
 
     def __str__(self):
diff --git a/pyteal/ast/gtxn_test.py b/pyteal/ast/gtxn_test.py
@@ -2,19 +2,57 @@
 
 from .. import *
 
+teal6Options = CompileOptions(version=6)
 
-def test_gtxn_invalid():
-    with pytest.raises(TealInputError):
-        Gtxn[-1].fee()
-
-    with pytest.raises(TealInputError):
-        Gtxn[MAX_GROUP_SIZE + 1].sender()
 
-    with pytest.raises(TealTypeError):
-        Gtxn[Pop(Int(0))].sender()
-
-    with pytest.raises(TealTypeError):
-        Gtxn[Bytes("index")].sender()
+def test_gtxn_invalid():
+    for f, e in [
+        (lambda: Gtxn[-1], TealInputError),
+        (lambda: Gtxn[MAX_GROUP_SIZE + 1], TealInputError),
+        (lambda: Gtxn[Pop(Int(0))], TealTypeError),
+        (lambda: Gtxn[Bytes("index")], TealTypeError),
+    ]:
+        with pytest.raises(e):
+            f()
+
+
+def test_gtxn_expr_invalid():
+    for f, e in [
+        (lambda: GtxnExpr(Assert(Int(1)), TxnField.sender), TealTypeError),
+    ]:
+        with pytest.raises(e):
+            f()
+
+
+def test_gtxn_expr_valid():
+    [
+        e.__teal__(teal6Options)
+        for e in [
+            GtxnExpr(1, TxnField.sender),
+            GtxnExpr(Int(1), TxnField.sender),
+        ]
+    ]
+
+
+def test_gtxna_expr_invalid():
+    for f, e in [
+        (lambda: GtxnaExpr("Invalid_type", TxnField.assets, 1), TealInputError),
+        (lambda: GtxnaExpr(1, TxnField.assets, "Invalid_type"), TealInputError),
+        (lambda: GtxnaExpr(Assert(Int(1)), TxnField.assets, 1), TealTypeError),
+        (lambda: GtxnaExpr(1, TxnField.assets, Assert(Int(1))), TealTypeError),
+    ]:
+        with pytest.raises(e):
+            f()
+
+
+def test_gtxna_expr_valid():
+    [
+        e.__teal__(teal6Options)
+        for e in [
+            GtxnaExpr(1, TxnField.assets, 1),
+            GtxnaExpr(Int(1), TxnField.assets, Int(1)),
+        ]
+    ]
 
 
 # txn_test.py performs additional testing
diff --git a/pyteal/ast/txn.py b/pyteal/ast/txn.py
@@ -161,6 +161,15 @@ def type_of(self):
 class TxnaExpr(LeafExpr):
     """An expression that accesses a transaction array field from the current transaction."""
 
+    @staticmethod
+    def __validate_index_or_throw(index: Union[int, Expr]):
+        if not isinstance(index, (int, Expr)):
+            raise TealInputError(
+                f"Invalid index type:  Expected int or Expr, but received {index}."
+            )
+        if isinstance(index, Expr):
+            require_type(index, TealType.uint64)
+
     def __init__(
         self,
         staticOp: Op,
@@ -172,6 +181,8 @@ def __init__(
         super().__init__()
         if not field.is_array:
             raise TealInputError("Unexpected non-array field: {}".format(field))
+        self.__validate_index_or_throw(index)
+
         self.staticOp = staticOp
         self.dynamicOp = dynamicOp
         self.name = name
