Alternate syntax for If-else statements (algorand#77)

* Add alternate if-else syntax with unit tests

* Fix typing errors for mypy

* Revise alternate syntax to have recursive structure and add more unit tests

* Change TealCompileError arg

pyteal/ast/if_.py

@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING
 
+from ..errors import TealCompileError, TealInputError
 from ..types import TealType, require_type, types_match
 from ..ir import TealSimpleBlock, TealConditionalBlock
 from .expr import Expr
@@ -10,7 +11,7 @@
 class If(Expr):
     """Simple two-way conditional expression."""
 
-    def __init__(self, cond: Expr, thenBranch: Expr, elseBranch: Expr = None) -> None:
+    def __init__(self, cond: Expr, thenBranch: Expr = None, 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
@@ -25,17 +26,26 @@ def __init__(self, cond: Expr, thenBranch: Expr, elseBranch: Expr = None) -> Non
         """
         super().__init__()
         require_type(cond.type_of(), TealType.uint64)
+        # Flag to denote and check whether the new If().Then() syntax is being used
+        self.alternateSyntaxFlag = False
 
-        if elseBranch is None:
-            require_type(thenBranch.type_of(), TealType.none)
+        if thenBranch:
+            if elseBranch:
+                require_type(thenBranch.type_of(), elseBranch.type_of())
+            else:
+                # If there is only a thenBranch, then it should evaluate to none type
+                require_type(thenBranch.type_of(), TealType.none)
         else:
-            require_type(thenBranch.type_of(), elseBranch.type_of())
+            self.alternateSyntaxFlag = True
 
         self.cond = cond
         self.thenBranch = thenBranch
         self.elseBranch = elseBranch
 
     def __teal__(self, options: 'CompileOptions'):
+        if self.thenBranch is None:
+            raise TealCompileError("If expression must have a thenBranch", self)
+
         condStart, condEnd = self.cond.__teal__(options)
         thenStart, thenEnd = self.thenBranch.__teal__(options)
         end = TealSimpleBlock([])
@@ -56,11 +66,51 @@ def __teal__(self, options: 'CompileOptions'):
         return condStart, end
 
     def __str__(self):
+        if self.thenBranch is None:
+            raise TealCompileError("If expression must have a thenBranch", 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):
+        if self.thenBranch is None:
+            raise TealCompileError("If expression must have a thenBranch", self) 
         return self.thenBranch.type_of()
 
+    def Then(self, thenBranch: Expr):
+        if not self.alternateSyntaxFlag:
+            raise TealInputError("Cannot mix two different If syntax styles")
+
+        if not self.elseBranch:
+            self.thenBranch = thenBranch
+        else:
+            if not isinstance(self.elseBranch, If):
+                raise TealInputError("Else-Then block is malformed")
+            self.elseBranch.Then(thenBranch)
+        return self
+
+    def ElseIf(self, cond):
+        if not self.alternateSyntaxFlag:
+            raise TealInputError("Cannot mix two different If syntax styles")
+
+        if not self.elseBranch:
+            self.elseBranch = If(cond)
+        else:
+            if not isinstance(self.elseBranch, If):
+                raise TealInputError("Else-ElseIf block is malformed")
+            self.elseBranch.ElseIf(cond)
+        return self
+
+    def Else(self, elseBranch: Expr):
+        if not self.alternateSyntaxFlag:
+            raise TealInputError("Cannot mix two different If syntax styles")
+
+        if not self.elseBranch:
+            self.elseBranch = elseBranch
+        else:
+            if not isinstance(self.elseBranch, If):
+                raise TealInputError("Else-Else block is malformed")
+            self.elseBranch.Else(elseBranch)
+        return self
+
 If.__module__ = "pyteal"

pyteal/ast/if_test.py

@@ -90,9 +90,151 @@ def test_if_invalid():
 
     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())
+
+    with pytest.raises(TealTypeError):
+        If(Int(0), Int(2))
+
+    with pytest.raises(TealCompileError):
+        expr = If(Int(0))
+        expr.__teal__(options)
+
+def test_if_alt_int():
+    args = [Int(0), Int(1), Int(2)]
+    expr = If(args[0]).Then(args[1]).Else(args[2])
+    assert expr.type_of() == TealType.uint64
+
+    expected, _ = args[0].__teal__(options)
+    thenBlock, _ = args[1].__teal__(options)
+    elseBlock, _ = args[2].__teal__(options)
+    expectedBranch = TealConditionalBlock([])
+    expectedBranch.setTrueBlock(thenBlock)
+    expectedBranch.setFalseBlock(elseBlock)
+    expected.setNextBlock(expectedBranch)
+    end = TealSimpleBlock([])
+    thenBlock.setNextBlock(end)
+    elseBlock.setNextBlock(end)
+
+    actual, _ = expr.__teal__(options)
+
+    assert actual == expected
+
+def test_if_alt_bytes():
+    args = [Int(1), Txn.sender(), Txn.receiver()]
+    expr = If(args[0]).Then(args[1]).Else(args[2])
+    assert expr.type_of() == TealType.bytes
+
+    expected, _ = args[0].__teal__(options)
+    thenBlock, _ = args[1].__teal__(options)
+    elseBlock, _ = args[2].__teal__(options)
+    expectedBranch = TealConditionalBlock([])
+    expectedBranch.setTrueBlock(thenBlock)
+    expectedBranch.setFalseBlock(elseBlock)
+    expected.setNextBlock(expectedBranch)
+    end = TealSimpleBlock([])
+    thenBlock.setNextBlock(end)
+    elseBlock.setNextBlock(end)
+
+    actual, _ = expr.__teal__(options)
+
+    assert actual == expected
+
+def test_if_alt_none():
+    args = [Int(0), Pop(Txn.sender()), Pop(Txn.receiver())]
+    expr = If(args[0]).Then(args[1]).Else(args[2])
+    assert expr.type_of() == TealType.none
+
+    expected, _ = args[0].__teal__(options)
+    thenBlockStart, thenBlockEnd = args[1].__teal__(options)
+    elseBlockStart, elseBlockEnd = args[2].__teal__(options)
+    expectedBranch = TealConditionalBlock([])
+    expectedBranch.setTrueBlock(thenBlockStart)
+    expectedBranch.setFalseBlock(elseBlockStart)
+    expected.setNextBlock(expectedBranch)
+    end = TealSimpleBlock([])
+    thenBlockEnd.setNextBlock(end)
+    elseBlockEnd.setNextBlock(end)
+
+    actual, _ = expr.__teal__(options)
+
+    assert actual == expected
+
+def test_elseif_syntax():
+    args = [Int(0), Int(1), Int(2), Int(3), Int(4)]
+    expr = If(args[0]).Then(args[1]).ElseIf(args[2]).Then(args[3]).Else(args[4])
+    assert expr.type_of() == TealType.uint64
+
+    elseExpr = If(args[2]).Then(args[3]).Else(args[4])
+    expected, _ = args[0].__teal__(options)
+    thenBlock, _ = args[1].__teal__(options)
+    elseStart, elseEnd = elseExpr.__teal__(options)
+    expectedBranch = TealConditionalBlock([])
+    expectedBranch.setTrueBlock(thenBlock)
+    expectedBranch.setFalseBlock(elseStart)
+    expected.setNextBlock(expectedBranch)
+    end = TealSimpleBlock([])
+    thenBlock.setNextBlock(end)
+    elseEnd.setNextBlock(end)
+
+    actual, _ = expr.__teal__(options)
+
+    assert actual == expected
+
+def test_elseif_multiple():
+    args = [Int(0), Int(1), Int(2), Int(3), Int(4), Int(5), Int(6)]
+    expr = If(args[0])\
+           .Then(args[1])\
+           .ElseIf(args[2])\
+           .Then(args[3])\
+           .ElseIf(args[4])\
+           .Then(args[5])\
+           .Else(args[6])
+    assert expr.type_of() == TealType.uint64
+
+    elseIfExpr = If(args[2], args[3], If(args[4], args[5], args[6]))
+    expected, _ = args[0].__teal__(options)
+    thenBlock, _ = args[1].__teal__(options)
+    elseStart, elseEnd = elseIfExpr.__teal__(options)
+    expectedBranch = TealConditionalBlock([])
+    expectedBranch.setTrueBlock(thenBlock)
+    expectedBranch.setFalseBlock(elseStart)
+    expected.setNextBlock(expectedBranch)
+    end = TealSimpleBlock([])
+    thenBlock.setNextBlock(end)
+    elseEnd.setNextBlock(end)
+
+    actual, _ = expr.__teal__(options)
+
+    assert actual == expected
+
+def test_if_invalid_alt_syntax():
+    with pytest.raises(TealCompileError):
+        expr = If(Int(0)).ElseIf(Int(1))
+        expr.__teal__(options)
+
+    with pytest.raises(TealCompileError):
+        expr = If(Int(0)).ElseIf(Int(1)).Then(Int(2))
+        expr.__teal__(options)
+
+    with pytest.raises(TealCompileError):
+        expr = If(Int(0)).Then(Int(1)).ElseIf(Int(2))
+        expr.__teal__(options)
+
+    with pytest.raises(TealCompileError):
+        expr = If(Int(0)).Then(Int(1)).ElseIf(Int(2))
+        expr.__teal__(options)
+
+    with pytest.raises(TealCompileError):
+        expr = If(Int(0)).Else(Int(1))
+        expr.__teal__(options)
+
+    with pytest.raises(TealInputError):
+        expr = If(Int(0)).Else(Int(1)).Then(Int(2))
+
+    with pytest.raises(TealInputError):
+        expr = If(Int(0)).Else(Int(1)).Else(Int(2))
+
+    with pytest.raises(TealInputError):
+        expr = If(Int(0), Pop(Int(1))).Else(Int(2))