Support logs, dynamic arg indexes, and other improvements

pyteal/ast/__init__.py

@@ -43,6 +43,7 @@
     MinBalance,
     BytesNot,
     BytesZero,
+    Log,
 )
 
 # binary ops
@@ -223,6 +224,7 @@
     "BytesGe",
     "BytesNot",
     "BytesZero",
+    "Log",
     "While",
     "For",
     "Break",

pyteal/ast/arg.py

@@ -1,8 +1,9 @@
-from typing import TYPE_CHECKING
+from typing import Union, cast, TYPE_CHECKING
 
-from ..types import TealType
+from ..types import TealType, require_type
 from ..ir import TealOp, Op, TealBlock
-from ..errors import TealInputError
+from ..errors import TealInputError, verifyTealVersion
+from .expr import Expr
 from .leafexpr import LeafExpr
 
 if TYPE_CHECKING:
@@ -12,28 +13,39 @@
 class Arg(LeafExpr):
     """An expression to get an argument when running in signature verification mode."""
 
-    def __init__(self, index: int) -> None:
+    def __init__(self, index: Union[int, Expr]) -> None:
         """Get an argument for this program.
 
         Should only be used in signature verification mode. For application mode arguments, see
         :any:`TxnObject.application_args`.
 
         Args:
-            index: The integer index of the argument to get. Must be between 0 and 255 inclusive.
+            index: The index of the argument to get. The index must be between 0 and 255 inclusive.
+                Starting in TEAL v5, the index may be a PyTeal expression that evaluates to uint64.
         """
         super().__init__()
 
-        if type(index) is not int:
-            raise TealInputError("invalid arg input type {}".format(type(index)))
-
-        if index < 0 or index > 255:
-            raise TealInputError("invalid arg index {}".format(index))
+        if type(index) is int:
+            if index < 0 or index > 255:
+                raise TealInputError("invalid arg index {}".format(index))
+        else:
+            require_type(cast(Expr, index).type_of(), TealType.uint64)
 
         self.index = index
 
     def __teal__(self, options: "CompileOptions"):
-        op = TealOp(self, Op.arg, self.index)
-        return TealBlock.FromOp(options, op)
+        if type(self.index) is int:
+            op = TealOp(self, Op.arg, self.index)
+            return TealBlock.FromOp(options, op)
+
+        verifyTealVersion(
+            Op.args.min_version,
+            options.version,
+            "TEAL version too low to use dynamic indexes with Arg",
+        )
+
+        op = TealOp(self, Op.args)
+        return TealBlock.FromOp(options, op, cast(Expr, self.index))
 
     def __str__(self):
         return "(arg {})".format(self.index)

pyteal/ast/arg_test.py

@@ -2,18 +2,47 @@
 
 from .. import *
 
+# this is not necessary but mypy complains if it's not included
+from .. import CompileOptions
 
-def test_arg():
-    expr = Arg(0)
+teal2Options = CompileOptions(version=2)
+teal4Options = CompileOptions(version=4)
+teal5Options = CompileOptions(version=5)
+
+
+def test_arg_static():
+    for i in range(256):
+        expr = Arg(i)
+        assert expr.type_of() == TealType.bytes
+        assert not expr.has_return()
+
+        expected = TealSimpleBlock([TealOp(expr, Op.arg, i)])
+
+        actual, _ = expr.__teal__(teal2Options)
+        assert actual == expected
+
+
+def test_arg_dynamic():
+    i = Int(7)
+    expr = Arg(i)
     assert expr.type_of() == TealType.bytes
-    expected = TealSimpleBlock([TealOp(expr, Op.arg, 0)])
-    actual, _ = expr.__teal__(CompileOptions())
+    assert not expr.has_return()
+
+    expected = TealSimpleBlock([TealOp(i, Op.int, 7), TealOp(expr, Op.args)])
+
+    actual, _ = expr.__teal__(teal5Options)
+    actual.addIncoming()
+    actual = TealBlock.NormalizeBlocks(actual)
+
     assert actual == expected
 
+    with pytest.raises(TealInputError):
+        expr.__teal__(teal4Options)
+
 
 def test_arg_invalid():
-    with pytest.raises(TealInputError):
-        Arg("k")
+    with pytest.raises(TealTypeError):
+        Arg(Bytes("k"))
 
     with pytest.raises(TealInputError):
         Arg(-1)

pyteal/ast/bytes.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import Union, cast, overload, TYPE_CHECKING
 
 from ..types import TealType, valid_base16, valid_base32, valid_base64
 from ..util import escapeStr
@@ -13,13 +13,23 @@
 class Bytes(LeafExpr):
     """An expression that represents a byte string."""
 
-    def __init__(self, *args: str) -> None:
+    @overload
+    def __init__(self, arg1: Union[str, bytes, bytearray]) -> None:
+        ...
+
+    @overload
+    def __init__(self, arg1: str, arg2: str) -> None:
+        ...
+
+    def __init__(self, arg1: Union[str, bytes, bytearray], arg2: str = None) -> None:
         """Create a new byte string.
 
         Depending on the encoding, there are different arguments to pass:
 
         For UTF-8 strings:
             Pass the string as the only argument. For example, ``Bytes("content")``.
+        For raw bytes or bytearray objects:
+            Pass the bytes or bytearray as the only argument. For example, ``Bytes(b"content")``.
         For base16, base32, or base64 strings:
             Pass the base as the first argument and the string as the second argument. For example,
             ``Bytes("base16", "636F6E74656E74")``, ``Bytes("base32", "ORFDPQ6ARJK")``,
@@ -29,35 +39,42 @@ def __init__(self, *args: str) -> None:
             ``Bytes("base16", "0x636F6E74656E74")``.
         """
         super().__init__()
-        if len(args) == 1:
-            self.base = "utf8"
-            self.byte_str = escapeStr(args[0])
-        elif len(args) == 2:
-            self.base, byte_str = args
+        if arg2 is None:
+            if type(arg1) is str:
+                self.base = "utf8"
+                self.byte_str = escapeStr(arg1)
+            elif type(arg1) in (bytes, bytearray):
+                self.base = "base16"
+                self.byte_str = cast(Union[bytes, bytearray], arg1).hex()
+            else:
+                raise TealInputError("Unknown argument type: {}".format(type(arg1)))
+        else:
+            if type(arg1) is not str:
+                raise TealInputError("Unknown type for base: {}".format(type(arg1)))
+
+            if type(arg2) is not str:
+                raise TealInputError("Unknown type for value: {}".format(type(arg2)))
+
+            self.base = arg1
+
             if self.base == "base32":
-                valid_base32(byte_str)
-                self.byte_str = byte_str
+                valid_base32(arg2)
+                self.byte_str = arg2
             elif self.base == "base64":
-                self.byte_str = byte_str
-                valid_base64(byte_str)
+                self.byte_str = arg2
+                valid_base64(self.byte_str)
             elif self.base == "base16":
-                if byte_str.startswith("0x"):
-                    self.byte_str = byte_str[2:]
+                if arg2.startswith("0x"):
+                    self.byte_str = arg2[2:]
                 else:
-                    self.byte_str = byte_str
+                    self.byte_str = arg2
                 valid_base16(self.byte_str)
             else:
                 raise TealInputError(
                     "invalid base {}, need to be base32, base64, or base16.".format(
                         self.base
                     )
                 )
-        else:
-            raise TealInputError(
-                "Only 1 or 2 arguments are expected for Bytes constructor, you provided {}".format(
-                    len(args)
-                )
-            )
 
     def __teal__(self, options: "CompileOptions"):
         if self.base == "utf8":

pyteal/ast/bytes_test.py

@@ -113,7 +113,31 @@ def test_bytes_utf8_empty():
     assert actual == expected
 
 
+def test_bytes_raw():
+    for value in (b"hello world", bytearray(b"hello world")):
+        expr = Bytes(value)
+        assert expr.type_of() == TealType.bytes
+        expected = TealSimpleBlock([TealOp(expr, Op.byte, "0x" + value.hex())])
+        actual, _ = expr.__teal__(options)
+        assert actual == expected
+
+
+def test_bytes_raw_empty():
+    for value in (b"", bytearray(b"")):
+        expr = Bytes(value)
+        assert expr.type_of() == TealType.bytes
+        expected = TealSimpleBlock([TealOp(expr, Op.byte, "0x")])
+        actual, _ = expr.__teal__(options)
+        assert actual == expected
+
+
 def test_bytes_invalid():
+    with pytest.raises(TealInputError):
+        Bytes("base16", b"FF")
+
+    with pytest.raises(TealInputError):
+        Bytes(b"base16", "FF")
+
     with pytest.raises(TealInputError):
         Bytes("base23", "")
 

pyteal/ast/naryexpr.py

@@ -2,7 +2,7 @@
 
 from ..types import TealType, require_type
 from ..errors import TealInputError
-from ..ir import TealOp, Op, TealSimpleBlock
+from ..ir import TealOp, Op, TealSimpleBlock, TealBlock
 from .expr import Expr
 
 if TYPE_CHECKING:
@@ -19,12 +19,12 @@ def __init__(
         self, op: Op, inputType: TealType, outputType: TealType, args: Sequence[Expr]
     ):
         super().__init__()
-        if len(args) < 2:
-            raise TealInputError("NaryExpr requires at least two children.")
+        if len(args) == 0:
+            raise TealInputError("NaryExpr requires at least one child")
         for arg in args:
             if not isinstance(arg, Expr):
                 raise TealInputError(
-                    "Argument is not a pyteal expression: {}".format(arg)
+                    "Argument is not a PyTeal expression: {}".format(arg)
                 )
             require_type(arg.type_of(), inputType)
         self.op = op
@@ -69,8 +69,8 @@ def And(*args: Expr) -> NaryExpr:
 
     Produces 1 if all arguments are nonzero. Otherwise produces 0.
 
-    All arguments must be PyTeal expressions that evaluate to uint64, and there must be at least two
-    arguments.
+    All arguments must be PyTeal expressions that evaluate to uint64, and there must be at least one
+    argument.
 
     Example:
         ``And(Txn.amount() == Int(500), Txn.fee() <= Int(10))``
@@ -83,8 +83,8 @@ def Or(*args: Expr) -> NaryExpr:
 
     Produces 1 if any argument is nonzero. Otherwise produces 0.
 
-    All arguments must be PyTeal expressions that evaluate to uint64, and there must be at least two
-    arguments.
+    All arguments must be PyTeal expressions that evaluate to uint64, and there must be at least one
+    argument.
     """
     return NaryExpr(Op.logic_or, TealType.uint64, TealType.uint64, args)
 
@@ -95,8 +95,8 @@ def Concat(*args: Expr) -> NaryExpr:
     Produces a new byte string consisting of the contents of each of the passed in byte strings
     joined together.
 
-    All arguments must be PyTeal expressions that evaluate to bytes, and there must be at least two
-    arguments.
+    All arguments must be PyTeal expressions that evaluate to bytes, and there must be at least one
+    argument.
 
     Example:
         ``Concat(Bytes("hello"), Bytes(" "), Bytes("world"))``