Comment as op (#410)

* add comment expression to allow comments in emitted TEAL

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Unreleased
 
+## Added
+* Add the ability to insert comments in TEAL source file with the Comment method ([#410](https://github.com/algorand/pyteal/pull/410))
+
 ## Fixed
 * Fix AST duplication bug in `String.set` when called with an `Expr` argument ([#508](https://github.com/algorand/pyteal/pull/508))
 

pyteal/__init__.pyi

@@ -82,6 +82,7 @@ __all__ = [
     "BytesXor",
     "BytesZero",
     "CallConfig",
+    "Comment",
     "CompileOptions",
     "Concat",
     "Cond",

pyteal/ast/__init__.py

@@ -42,6 +42,7 @@
 from pyteal.ast.tmpl import Tmpl
 from pyteal.ast.nonce import Nonce
 from pyteal.ast.pragma import Pragma
+from pyteal.ast.comment import Comment
 
 # unary ops
 from pyteal.ast.unaryexpr import (
@@ -216,6 +217,7 @@
     "Tmpl",
     "Nonce",
     "Pragma",
+    "Comment",
     "UnaryExpr",
     "Btoi",
     "Itob",

pyteal/ast/comment.py

@@ -0,0 +1,66 @@
+from typing import TYPE_CHECKING, Tuple
+
+from pyteal.errors import TealInputError
+from pyteal.types import TealType
+from pyteal.ir import TealBlock, TealSimpleBlock, TealOp, Op
+from pyteal.ast.expr import Expr
+from pyteal.ast.seq import Seq
+
+if TYPE_CHECKING:
+    from pyteal.compiler import CompileOptions
+
+
+class CommentExpr(Expr):
+    """Represents a single line comment in TEAL source code.
+
+    This class is intentionally hidden because it's too basic to directly expose. Anything exposed
+    to users should be able to handle multi-line comments by breaking them apart and using this
+    class.
+    """
+
+    def __init__(self, single_line_comment: str) -> None:
+        super().__init__()
+        if "\n" in single_line_comment or "\r" in single_line_comment:
+            raise TealInputError(
+                "Newlines should not be present in the CommentExpr constructor"
+            )
+        self.comment = single_line_comment
+
+    def __teal__(self, options: "CompileOptions") -> Tuple[TealBlock, TealSimpleBlock]:
+        op = TealOp(self, Op.comment, self.comment)
+        return TealBlock.FromOp(options, op)
+
+    def __str__(self):
+        return f'(Comment "{self.comment}")'
+
+    def type_of(self):
+        return TealType.none
+
+    def has_return(self):
+        return False
+
+
+CommentExpr.__module__ = "pyteal"
+
+
+def Comment(comment: str, expr: Expr = None) -> Expr:
+    """Wrap an existing expression with a comment.
+
+    This comment will be present in the compiled TEAL source immediately before the first op of the
+    expression.
+
+    Note that when TEAL source is assembled into bytes, all comments are omitted.
+
+    Args:
+        comment: The comment that will be associated with the expression.
+        expr: The expression to be commented.
+
+    Returns:
+        A new expression which is functionally equivalent to the input expression, but which will
+        compile with the given comment string.
+    """
+    lines = comment.splitlines()
+    comment_lines: list[Expr] = [CommentExpr(line) for line in lines]
+    if expr is not None:
+        comment_lines.append(expr)
+    return Seq(*comment_lines)

pyteal/ast/comment_test.py

@@ -0,0 +1,116 @@
+import pytest
+
+import pyteal as pt
+from pyteal.ast.comment import CommentExpr
+
+options = pt.CompileOptions()
+
+
+def test_CommentExpr():
+    for comment in ("", "hello world", " // a b c //    \t . "):
+        expr = CommentExpr(comment)
+        assert expr.comment == comment
+        assert expr.type_of() == pt.TealType.none
+        assert expr.has_return() is False
+
+        expected = pt.TealSimpleBlock(
+            [
+                pt.TealOp(expr, pt.Op.comment, comment),
+            ]
+        )
+
+        actual, _ = expr.__teal__(options)
+        actual.addIncoming()
+        actual = pt.TealBlock.NormalizeBlocks(actual)
+
+        assert actual == expected
+
+    for newline in ("\n", "\r\n", "\r"):
+        with pytest.raises(
+            pt.TealInputError,
+            match=r"Newlines should not be present in the CommentExpr constructor$",
+        ):
+            CommentExpr(f"one line{newline}two lines")
+
+
+def test_Comment_Expr_empty():
+    comment = "dope"
+    expr = pt.Comment(comment)
+    assert type(expr) is pt.Seq
+    assert len(expr.args) == 1
+
+    assert getattr(expr.args[0], "comment") == comment
+
+
+def test_Comment_empty():
+    to_wrap = pt.Int(1)
+    comment = ""
+    expr = pt.Comment(comment, to_wrap)
+    assert type(expr) is pt.Seq
+    assert len(expr.args) == 1
+
+    assert expr.args[0] is to_wrap
+
+
+def test_Comment_single_line():
+    to_wrap = pt.Int(1)
+    comment = "just an int"
+    expr = pt.Comment(comment, to_wrap)
+    assert type(expr) is pt.Seq
+    assert len(expr.args) == 2
+
+    assert type(expr.args[0]) is CommentExpr
+    assert expr.args[0].comment == comment
+
+    assert expr.args[1] is to_wrap
+
+    version = 6
+    expected_teal = f"""#pragma version {version}
+// {comment}
+int 1
+return"""
+    actual_teal = pt.compileTeal(
+        pt.Return(expr), version=version, mode=pt.Mode.Application
+    )
+    assert actual_teal == expected_teal
+
+
+def test_Comment_multi_line():
+    to_wrap = pt.Int(1)
+    comment = """just an int
+but its really more than that isnt it? an integer here is a uint64 stack type but looking further what does that mean? 
+You might say its a 64 bit representation of an element of the set Z and comes from the latin `integer` meaning `whole`
+since it has no fractional part. You might also say this run on comment has gone too far. See https://en.wikipedia.org/wiki/Integer for more details 
+"""
+
+    comment_parts = [
+        "just an int",
+        "but its really more than that isnt it? an integer here is a uint64 stack type but looking further what does that mean? ",
+        "You might say its a 64 bit representation of an element of the set Z and comes from the latin `integer` meaning `whole`",
+        "since it has no fractional part. You might also say this run on comment has gone too far. See https://en.wikipedia.org/wiki/Integer for more details ",
+    ]
+
+    expr = pt.Comment(comment, to_wrap)
+    assert type(expr) is pt.Seq
+    assert len(expr.args) == 5
+
+    for i, part in enumerate(comment_parts):
+        arg = expr.args[i]
+        assert type(arg) is CommentExpr
+        assert arg.comment == part
+
+    assert expr.args[4] is to_wrap
+
+    version = 6
+    expected_teal = f"""#pragma version {version}
+// just an int
+// but its really more than that isnt it? an integer here is a uint64 stack type but looking further what does that mean? 
+// You might say its a 64 bit representation of an element of the set Z and comes from the latin `integer` meaning `whole`
+// since it has no fractional part. You might also say this run on comment has gone too far. See https://en.wikipedia.org/wiki/Integer for more details 
+int 1
+return"""
+    actual_teal = pt.compileTeal(
+        pt.Return(expr), version=version, mode=pt.Mode.Application
+    )
+    print(actual_teal)
+    assert actual_teal == expected_teal

pyteal/ir/ops.py

@@ -31,6 +31,9 @@ def min_version(self) -> int:
         return self.value.min_version
 
     # fmt: off
+    # meta
+    comment             = OpType("//",                  Mode.Signature | Mode.Application, 0)
+    # avm 
     err                 = OpType("err",                 Mode.Signature | Mode.Application, 2)
     sha256              = OpType("sha256",              Mode.Signature | Mode.Application, 2)
     keccak256           = OpType("keccak256",           Mode.Signature | Mode.Application, 2)