Allow single branch if statements (#19)

pyteal/ast/if_.py

@@ -6,44 +6,57 @@
 class If(Expr):
     """Simple two-way conditional expression."""
 
-    def __init__(self, cond: Expr, thenBranch: Expr, elseBranch: Expr) -> None:
+    def __init__(self, cond: Expr, thenBranch: Expr, elseBranch: Expr = None) -> None:
         """Create a new If expression.
 
         When this If expression is executed, the condition will be evaluated, and if it produces a
         true value, thenBranch will be executed and used as the return value for this expression.
-        Otherwise, elseBranch will be executed and used as the return value.
+        Otherwise, elseBranch will be executed and used as the return value, if it is provided.
 
         Args:
             cond: The condition to check. Must evaluate to uint64.
             thenBranch: Expression to evaluate if the condition is true.
-            elseBranch: Expression to evaluate if the condition is false. Must evaluate to the same
-            type as thenBranch.
+            elseBranch (optional): Expression to evaluate if the condition is false. Must evaluate
+            to the same type as thenBranch, if provided. Defaults to None.
         """
         require_type(cond.type_of(), TealType.uint64)
-        require_type(thenBranch.type_of(), elseBranch.type_of())
 
+        if elseBranch is None:
+            require_type(thenBranch.type_of(), TealType.none)
+        else:
+            require_type(thenBranch.type_of(), elseBranch.type_of())
+
         self.cond = cond
         self.thenBranch = thenBranch
         self.elseBranch = elseBranch
 
     def __teal__(self):
-        cond = self.cond.__teal__()
-        l1 = new_label()
-        t_branch = self.thenBranch.__teal__()
-        e_branch = self.elseBranch.__teal__()
-        l2 = new_label()
-
-        teal = cond
-        teal.append(TealOp(Op.bnz, l1))
-        teal += e_branch
-        teal.append(TealOp(Op.b, l2))
-        teal.append(TealLabel(l1))
-        teal += t_branch
-        teal.append(TealLabel(l2))
+        teal = self.cond.__teal__()
+        end = new_label()
+        tealThen = self.thenBranch.__teal__()
+        tealElse = [] if self.elseBranch is None else self.elseBranch.__teal__()
+
+        if self.elseBranch is None:
+            teal.append(TealOp(Op.bz, end))
+            teal += tealThen
+        else:
+            # doing this swap so that labels remain consistent with previous If implementation.
+            then = end
+            end = new_label()
+
+            teal.append(TealOp(Op.bnz, then))
+            teal += tealElse
+            teal.append(TealOp(Op.b, end))
+            teal.append(TealLabel(then))
+            teal += tealThen
+
+        teal.append(TealLabel(end))
 
         return teal
 
     def __str__(self):
+        if self.elseBranch is None:
+            return "(If {} {})".format(self.cond, self.thenBranch)
         return "(If {} {} {})".format(self.cond, self.thenBranch, self.elseBranch)
 
     def type_of(self):

pyteal/ast/if_test.py

@@ -38,9 +38,51 @@ def test_if_bytes():
         labels[1]
     ]
 
+def test_if_none():
+    expr = If(Int(0), Pop(Txn.sender()), Pop(Txn.receiver()))
+    assert expr.type_of() == TealType.none
+    teal = expr.__teal__()
+    assert len(teal) == 9
+    labels = [teal[5], teal[8]]
+    assert all(isinstance(label, TealLabel) for label in labels)
+    assert len(labels) == len(set(labels))
+    assert teal == [
+        TealOp(Op.int, 0),
+        TealOp(Op.bnz, labels[0].label),
+        TealOp(Op.txn, "Receiver"),
+        TealOp(Op.pop),
+        TealOp(Op.b, labels[1].label),
+        labels[0],
+        TealOp(Op.txn, "Sender"),
+        TealOp(Op.pop),
+        labels[1]
+    ]
+
+def test_if_single():
+    expr = If(Int(0), Pop(Int(1)))
+    assert expr.type_of() == TealType.none
+    teal = expr.__teal__()
+    assert len(teal) == 5
+    labels = [teal[4]]
+    assert all(isinstance(label, TealLabel) for label in labels)
+    assert len(labels) == len(set(labels))
+    assert teal == [
+        TealOp(Op.int, 0),
+        TealOp(Op.bz, labels[0].label),
+        TealOp(Op.int, 1),
+        TealOp(Op.pop),
+        labels[0]
+    ]
+
 def test_if_invalid():
     with pytest.raises(TealTypeError):
         If(Int(0), Txn.amount(), Txn.sender())
 
     with pytest.raises(TealTypeError):
         If(Txn.sender(), Int(1), Int(0))
+
+    with pytest.raises(TealTypeError):
+        If(Int(0), Int(1))
+
+    with pytest.raises(TealTypeError):
+        If(Int(0), Txn.sender())