diff --git a/pyteal/ast/__init__.py b/pyteal/ast/__init__.py index 400228463..a131142dd 100644 --- a/pyteal/ast/__init__.py +++ b/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", diff --git a/pyteal/ast/arg.py b/pyteal/ast/arg.py index 301c1514b..2694dafe1 100644 --- a/pyteal/ast/arg.py +++ b/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) diff --git a/pyteal/ast/arg_test.py b/pyteal/ast/arg_test.py index 741065112..7d622ddb7 100644 --- a/pyteal/ast/arg_test.py +++ b/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) diff --git a/pyteal/ast/bytes.py b/pyteal/ast/bytes.py index 369c47363..1bb58b4dc 100644 --- a/pyteal/ast/bytes.py +++ b/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,22 +39,35 @@ 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( @@ -52,12 +75,6 @@ def __init__(self, *args: str) -> None: 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": diff --git a/pyteal/ast/bytes_test.py b/pyteal/ast/bytes_test.py index 1b479e3e3..3f0bb49f4 100644 --- a/pyteal/ast/bytes_test.py +++ b/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", "") diff --git a/pyteal/ast/naryexpr.py b/pyteal/ast/naryexpr.py index b15d173ca..97ab5b061 100644 --- a/pyteal/ast/naryexpr.py +++ b/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"))`` diff --git a/pyteal/ast/naryexpr_test.py b/pyteal/ast/naryexpr_test.py index 7be86aa0e..674c6d119 100644 --- a/pyteal/ast/naryexpr_test.py +++ b/pyteal/ast/naryexpr_test.py @@ -8,6 +8,18 @@ options = CompileOptions() +def test_and_one(): + arg = Int(1) + expr = And(arg) + assert expr.type_of() == TealType.uint64 + + expected = TealSimpleBlock([TealOp(arg, Op.int, 1)]) + + actual, _ = expr.__teal__(options) + + assert actual == expected + + def test_and_two(): args = [Int(1), Int(2)] expr = And(args[0], args[1]) @@ -74,9 +86,6 @@ def test_and_invalid(): with pytest.raises(TealInputError): And() - with pytest.raises(TealInputError): - And(Int(1)) - with pytest.raises(TealTypeError): And(Int(1), Txn.receiver()) @@ -87,6 +96,18 @@ def test_and_invalid(): And(Txn.receiver(), Txn.receiver()) +def test_or_one(): + arg = Int(1) + expr = Or(arg) + assert expr.type_of() == TealType.uint64 + + expected = TealSimpleBlock([TealOp(arg, Op.int, 1)]) + + actual, _ = expr.__teal__(options) + + assert actual == expected + + def test_or_two(): args = [Int(1), Int(0)] expr = Or(args[0], args[1]) @@ -153,9 +174,6 @@ def test_or_invalid(): with pytest.raises(TealInputError): Or() - with pytest.raises(TealInputError): - Or(Int(1)) - with pytest.raises(TealTypeError): Or(Int(1), Txn.receiver()) @@ -164,3 +182,71 @@ def test_or_invalid(): with pytest.raises(TealTypeError): Or(Txn.receiver(), Txn.receiver()) + + +def test_concat_one(): + arg = Bytes("a") + expr = Concat(arg) + assert expr.type_of() == TealType.bytes + + expected = TealSimpleBlock([TealOp(arg, Op.byte, '"a"')]) + + actual, _ = expr.__teal__(options) + + assert actual == expected + + +def test_concat_two(): + args = [Bytes("a"), Bytes("b")] + expr = Concat(args[0], args[1]) + assert expr.type_of() == TealType.bytes + + expected = TealSimpleBlock( + [ + TealOp(args[0], Op.byte, '"a"'), + TealOp(args[1], Op.byte, '"b"'), + TealOp(expr, Op.concat), + ] + ) + + actual, _ = expr.__teal__(options) + actual.addIncoming() + actual = TealBlock.NormalizeBlocks(actual) + + assert actual == expected + + +def test_concat_three(): + args = [Bytes("a"), Bytes("b"), Bytes("c")] + expr = Concat(args[0], args[1], args[2]) + assert expr.type_of() == TealType.bytes + + expected = TealSimpleBlock( + [ + TealOp(args[0], Op.byte, '"a"'), + TealOp(args[1], Op.byte, '"b"'), + TealOp(expr, Op.concat), + TealOp(args[2], Op.byte, '"c"'), + TealOp(expr, Op.concat), + ] + ) + + actual, _ = expr.__teal__(options) + actual.addIncoming() + actual = TealBlock.NormalizeBlocks(actual) + + assert actual == expected + + +def test_concat_invalid(): + with pytest.raises(TealInputError): + Concat() + + with pytest.raises(TealTypeError): + Concat(Int(1), Txn.receiver()) + + with pytest.raises(TealTypeError): + Concat(Txn.receiver(), Int(1)) + + with pytest.raises(TealTypeError): + Concat(Int(1), Int(2)) diff --git a/pyteal/ast/unaryexpr.py b/pyteal/ast/unaryexpr.py index 15f77d947..47647f4d9 100644 --- a/pyteal/ast/unaryexpr.py +++ b/pyteal/ast/unaryexpr.py @@ -1,6 +1,7 @@ from typing import TYPE_CHECKING from ..types import TealType, require_type +from ..errors import verifyTealVersion from ..ir import TealOp, Op, TealBlock from .expr import Expr @@ -21,6 +22,12 @@ def __init__( self.arg = arg def __teal__(self, options: "CompileOptions"): + verifyTealVersion( + self.op.min_version, + options.version, + "TEAL version too low to use op {}".format(self.op), + ) + return TealBlock.FromOp(options, TealOp(self, self.op), self.arg) def __str__(self): @@ -154,3 +161,17 @@ def BytesZero(arg: Expr) -> UnaryExpr: Requires TEAL version 4 or higher. """ return UnaryExpr(Op.bzero, TealType.uint64, TealType.bytes, arg) + + +def Log(message: Expr) -> UnaryExpr: + """Write a message to log state of the current application. + + This will fail if called more than :code:`MaxLogCalls` times in a program (32 as of TEAL v5), or + if the sum of the lengths of all logged messages in a program exceeds 1024 bytes. + + Args: + message: The message to write. Must evaluate to bytes. + + Requires TEAL version 5 or higher. + """ + return UnaryExpr(Op.log, TealType.bytes, TealType.none, message) diff --git a/pyteal/ast/unaryexpr_test.py b/pyteal/ast/unaryexpr_test.py index 8ff633286..7da54a3a4 100644 --- a/pyteal/ast/unaryexpr_test.py +++ b/pyteal/ast/unaryexpr_test.py @@ -8,6 +8,7 @@ teal2Options = CompileOptions(version=2) teal3Options = CompileOptions(version=3) teal4Options = CompileOptions(version=4) +teal5Options = CompileOptions(version=5) def test_btoi(): @@ -373,3 +374,28 @@ def test_b_zero(): def test_b_zero_invalid(): with pytest.raises(TealTypeError): BytesZero(Bytes("base16", "0x11")) + + +def test_log(): + arg = Bytes("message") + expr = Log(arg) + assert expr.type_of() == TealType.none + assert not expr.has_return() + + expected = TealSimpleBlock( + [TealOp(arg, Op.byte, '"message"'), TealOp(expr, Op.log)] + ) + + actual, _ = expr.__teal__(teal5Options) + actual.addIncoming() + actual = TealBlock.NormalizeBlocks(actual) + + assert actual == expected + + with pytest.raises(TealInputError): + expr.__teal__(teal4Options) + + +def test_log_invalid(): + with pytest.raises(TealTypeError): + Log(Int(7))