Add opcode support for itxnas and gitxnas

pyteal/ast/gitxn.py

@@ -1,9 +1,10 @@
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, cast, Union
 
 from pyteal.config import MAX_GROUP_SIZE
 
 from ..errors import TealInputError, verifyFieldVersion, verifyTealVersion
 from ..ir import TealOp, Op, TealBlock
+from .expr import Expr
 from .txn import TxnExpr, TxnField, TxnObject, TxnaExpr
 
 if TYPE_CHECKING:
@@ -46,8 +47,8 @@ def __teal__(self, options: "CompileOptions"):
 class GitxnaExpr(TxnaExpr):
     """An expression that accesses an inner transaction array field from an inner transaction in the last inner group."""
 
-    def __init__(self, txnIndex: int, field: TxnField, index: int) -> None:
-        super().__init__(Op.gitxna, None, "Gitxna", field, index)
+    def __init__(self, txnIndex: int, field: TxnField, index: Union[int, Expr]) -> None:
+        super().__init__(Op.gitxna, Op.gitxnas, "Gitxna", field, index)
         self.txnIndex = txnIndex
 
     def __str__(self):
@@ -57,18 +58,23 @@ 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 type(self.index) is not int:
-            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:
+            opToUse = Op.gitxnas
 
         verifyTealVersion(
-            Op.gitxna.min_version, options.version, "TEAL version too low to use gitxna"
+            opToUse.min_version,
+            options.version,
+            "TEAL version too low to use op {}".format(opToUse),
         )
-        op = TealOp(self, Op.gitxna, self.txnIndex, self.field.arg_name, self.index)
-        return TealBlock.FromOp(options, op)
+
+        if type(self.index) is int:
+            op = TealOp(self, opToUse, self.txnIndex, self.field.arg_name, self.index)
+            return TealBlock.FromOp(options, op)
+        op = TealOp(self, opToUse, self.txnIndex, self.field.arg_name)
+        return TealBlock.FromOp(options, op, cast(Expr, self.index))
 
 
 GitxnaExpr.__module__ = "pyteal"
@@ -85,7 +91,7 @@ def __getitem__(self, txnIndex: int) -> TxnObject:
 
         if txnIndex < 0 or txnIndex >= MAX_GROUP_SIZE:
             raise TealInputError(
-                "Invalid Gtxn index {}, shoud be in [0, {})".format(
+                "Invalid Gtxn index {}, should be in [0, {})".format(
                     txnIndex, MAX_GROUP_SIZE
                 )
             )

pyteal/ast/itxn.py

@@ -120,7 +120,7 @@ def Submit(cls) -> Expr:
         :any:`InnerTxnBuilder.Begin` and :any:`InnerTxnBuilder.SetField` must be called before
         submitting an inner transaction.
 
-        This will fail fail if 16 inner transactions have already been executed, or if the
+        This will fail if 256 inner transactions have already been executed, or if the
         inner transaction itself fails. Upon failure, the current program will immediately exit and
         fail as well.
 
@@ -204,7 +204,8 @@ def SetFields(cls, fields: Dict[TxnField, Union[Expr, List[Expr]]]) -> Expr:
 InnerTxnBuilder.__module__ = "pyteal"
 
 InnerTxn: TxnObject = TxnObject(
-    TxnExprBuilder(Op.itxn, "InnerTxn"), TxnaExprBuilder(Op.itxna, None, "InnerTxna")
+    TxnExprBuilder(Op.itxn, "InnerTxn"),
+    TxnaExprBuilder(Op.itxna, Op.itxnas, "InnerTxna"),
 )
 
 InnerTxn.__module__ = "pyteal"

pyteal/ast/txn_test.py

@@ -96,9 +96,9 @@ def test_txn_fields():
             [],
             [TealOp(dynamicGtxnArg, Op.int, 0)],
         ),
-        (InnerTxn, Op.itxn, Op.itxna, None, [], []),
+        (InnerTxn, Op.itxn, Op.itxna, Op.itxnas, [], []),
         *[
-            (Gitxn[i], Op.gitxn, Op.gitxna, None, [i], [])
+            (Gitxn[i], Op.gitxn, Op.gitxna, Op.gitxnas, [i], [])
             for i in range(MAX_GROUP_SIZE)
         ],
     ]