Fix If type bug (#329)

* Fix If type bug * Fix multi test
algorand · May 9, 2022 · 22deba5 · 22deba5
1 parent f97e02e
commit 22deba5
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/pyteal/ast/if_.py b/pyteal/ast/if_.py
@@ -1,6 +1,9 @@
 from typing import TYPE_CHECKING
 
-from pyteal.errors import TealCompileError, TealInputError
+from pyteal.errors import (
+    TealCompileError,
+    TealInputError,
+)
 from pyteal.types import TealType, require_type
 from pyteal.ir import TealSimpleBlock, TealConditionalBlock
 from pyteal.ast.expr import Expr
@@ -78,6 +81,11 @@ def __str__(self):
     def type_of(self):
         if self.thenBranch is None:
             raise TealCompileError("If expression must have a thenBranch", self)
+
+        if self.elseBranch is None:
+            # if there is only a thenBranch, it must evaluate to TealType.none
+            require_type(self.thenBranch, TealType.none)
+
         return self.thenBranch.type_of()
 
     def has_return(self):
@@ -119,7 +127,11 @@ def Else(self, elseBranch: Expr):
         if not self.alternateSyntaxFlag:
             raise TealInputError("Cannot mix two different If syntax styles")
 
+        if not self.thenBranch:
+            raise TealInputError("Must set Then branch before Else branch")
+
         if not self.elseBranch:
+            require_type(elseBranch, self.thenBranch.type_of())
             self.elseBranch = elseBranch
         else:
             if not isinstance(self.elseBranch, If):

diff --git a/pyteal/ast/if_test.py b/pyteal/ast/if_test.py
@@ -252,7 +252,14 @@ def test_if_invalid_alt_syntax():
         expr = pt.If(pt.Int(0)).Then(pt.Int(1)).ElseIf(pt.Int(2))
         expr.__teal__(options)
 
-    with pytest.raises(pt.TealCompileError):
+    with pytest.raises(pt.TealTypeError):
+        expr = pt.If(pt.Int(0)).Then(pt.Int(2))
+        expr.type_of()
+
+    with pytest.raises(pt.TealInputError):
+        pt.If(pt.Int(0)).Else(pt.Int(1)).Then(pt.Int(2))
+
+    with pytest.raises(pt.TealInputError):
         expr = pt.If(pt.Int(0)).Else(pt.Int(1))
         expr.__teal__(options)
 

diff --git a/pyteal/ast/multi_test.py b/pyteal/ast/multi_test.py
@@ -37,7 +37,7 @@ def __test_single_conditional(
     ifExpr = (
         pt.If(expr.output_slots[1].load())
         .Then(expr.output_slots[0].load())
-        .Else(pt.Bytes("None"))
+        .Else(pt.App.globalGet(pt.Bytes("None")))
     )
     ifBlockStart, _ = ifExpr.__teal__(options)
 
@@ -165,7 +165,7 @@ def test_multi_value():
                     reducer = (
                         lambda value, hasValue: pt.If(hasValue)
                         .Then(value)
-                        .Else(pt.Bytes("None"))
+                        .Else(pt.App.globalGet(pt.Bytes("None")))
                     )
                     expr = pt.MultiValue(
                         op, [type, pt.TealType.uint64], immediate_args=iargs, args=args