adding comment kwarg to assert expression (#510)

* adding comment kwarg to assert expression

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Added
 * Add the ability to insert comments in TEAL source file with the Comment method ([#410](https://github.com/algorand/pyteal/pull/410))
 * Static and Dynamic Byte Array convenience classes ([#500](https://github.com/algorand/pyteal/pull/500))
+* Add a `comment` keyword argument to the Assert expression that will place the comment immediately above the `assert` op in the resulting TEAL ([#510](https://github.com/algorand/pyteal/pull/510))
 
 ## Fixed
 * Fix AST duplication bug in `String.set` when called with an `Expr` argument ([#508](https://github.com/algorand/pyteal/pull/508))

pyteal/ast/assert_.py

@@ -3,6 +3,7 @@
 from pyteal.types import TealType, require_type
 from pyteal.ir import TealOp, Op, TealBlock, TealSimpleBlock, TealConditionalBlock
 from pyteal.ast.expr import Expr
+from pyteal.ast.comment import Comment
 from pyteal.ast.seq import Seq
 
 if TYPE_CHECKING:
@@ -12,27 +13,39 @@
 class Assert(Expr):
     """A control flow expression to verify that a condition is true."""
 
-    def __init__(self, cond: Expr, *additional_conds: Expr) -> None:
+    def __init__(
+        self, cond: Expr, *additional_conds: Expr, comment: str = None
+    ) -> None:
         """Create an assert statement that raises an error if the condition is false.
 
         Args:
             cond: The condition to check. Must evaluate to a uint64.
             *additional_conds: Additional conditions to check. Must evaluate to uint64.
+            comment: String comment to place on the line immediately prior to the assert op
         """
         super().__init__()
         require_type(cond, TealType.uint64)
         for cond_single in additional_conds:
             require_type(cond_single, TealType.uint64)
+
+        self.comment = comment
         self.cond = [cond] + list(additional_conds)
 
     def __teal__(self, options: "CompileOptions"):
         if len(self.cond) > 1:
-            asserts = [Assert(cond) for cond in self.cond]
+            asserts: list[Expr] = []
+            for cond in self.cond:
+                asrt = Assert(cond, comment=self.comment)
+                asrt.trace = cond.trace
+                asserts.append(asrt)
             return Seq(*asserts).__teal__(options)
 
         if options.version >= Op.assert_.min_version:
             # use assert op if available
-            return TealBlock.FromOp(options, TealOp(self, Op.assert_), self.cond[0])
+            conds: list[Expr] = [self.cond[0]]
+            if self.comment is not None:
+                conds.append(Comment(self.comment))
+            return TealBlock.FromOp(options, TealOp(self, Op.assert_), *conds)
 
         # if assert op is not available, use branches and err
         condStart, condEnd = self.cond[0].__teal__(options)

pyteal/ast/assert_test.py

@@ -73,6 +73,54 @@ def test_teal_3_assert_multi():
         assert actual == expected
 
 
+def test_assert_comment():
+    comment = "Make sure 1 is true"
+    expr = pt.Assert(pt.Int(1), comment=comment)
+    assert expr.type_of() == pt.TealType.none
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(None, pt.Op.int, 1),
+            pt.TealOp(None, pt.Op.comment, comment),
+            pt.TealOp(None, pt.Op.assert_),
+        ]
+    )
+
+    actual, _ = expr.__teal__(avm3Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    with pt.TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+
+def test_assert_comment_multi():
+    comment = "Make sure numbers > 0 are true"
+    expr = pt.Assert(pt.Int(1), pt.Int(2), pt.Int(3), comment=comment)
+    assert expr.type_of() == pt.TealType.none
+
+    expected = pt.TealSimpleBlock(
+        [
+            pt.TealOp(None, pt.Op.int, 1),
+            pt.TealOp(None, pt.Op.comment, comment),
+            pt.TealOp(None, pt.Op.assert_),
+            pt.TealOp(None, pt.Op.int, 2),
+            pt.TealOp(None, pt.Op.comment, comment),
+            pt.TealOp(None, pt.Op.assert_),
+            pt.TealOp(None, pt.Op.int, 3),
+            pt.TealOp(None, pt.Op.comment, comment),
+            pt.TealOp(None, pt.Op.assert_),
+        ]
+    )
+
+    actual, _ = expr.__teal__(avm3Options)
+    actual.addIncoming()
+    actual = pt.TealBlock.NormalizeBlocks(actual)
+
+    with pt.TealComponent.Context.ignoreExprEquality():
+        assert actual == expected
+
+
 def test_assert_invalid():
     with pytest.raises(pt.TealTypeError):
         pt.Assert(pt.Txn.receiver())

pyteal/ast/comment_test.py

@@ -112,5 +112,5 @@ def test_Comment_multi_line():
     actual_teal = pt.compileTeal(
         pt.Return(expr), version=version, mode=pt.Mode.Application
     )
-    print(actual_teal)
+
     assert actual_teal == expected_teal

pyteal/ast/cond_test.py

@@ -169,8 +169,6 @@ def test_cond_two_pred_multi():
     expected = cond1
 
     actual, _ = expr.__teal__(options)
-    print(actual)
-    print(expected)
 
     with pt.TealComponent.Context.ignoreExprEquality():
         assert actual == expected