Add initial implementation for forelse/whileelse

nottheswimmer · Jul 6, 2021 · 0242fa8 · 0242fa8
1 parent cd12ebb
commit 0242fa8
Show file tree

Hide file tree

Showing 8 changed files with 250 additions and 101 deletions.
diff --git a/README.md b/README.md
@@ -4219,6 +4219,157 @@ func main() {
 	fmt.Println(x, y, z, a, b, c, d)
 }
 ```
+### scope
+#### Python
+```python
+import random
+
+
+def main():
+    if random.random() > 0.5:
+        a = 1
+    else:
+        a = 2
+
+    if random.random() > 0.5:
+        if random.random() > 0.5:
+            b = 1
+        else:
+            b = 2
+    else:
+        b = 3
+
+    def hello_world():
+        c = 3
+        print(c)
+
+    hello_world()
+    print(a, b)
+
+
+if __name__ == '__main__':
+    main()
+```
+#### Go
+```go
+package main
+
+import (
+	"fmt"
+	"math/rand"
+	"time"
+)
+
+func init() {
+	rand.Seed(time.Now().UnixNano())
+}
+
+func main() {
+	var a int
+	var b int
+	if rand.Float64() > 0.5 {
+		a = 1
+	} else {
+		a = 2
+	}
+	if rand.Float64() > 0.5 {
+		if rand.Float64() > 0.5 {
+			b = 1
+		} else {
+			b = 2
+		}
+	} else {
+		b = 3
+	}
+	hello_world := func() {
+		c := 3
+		fmt.Println(c)
+	}
+	hello_world()
+	fmt.Println(a, b)
+}
+```
+### forelse
+#### Python
+```python
+def main():
+    for x in range(4):
+        if x == 5:
+            break
+    else:
+        print("Well of course that didn't happen")
+
+    for x in range(7):
+        if x == 5:
+            break
+    else:
+        print("H-hey wait!")
+
+    i = 0
+    while i < 3:
+        print("Works with while too")
+        for x in range(3):
+            print("BTW don't worry about nested breaks")
+            break
+        if i == 10:
+            break
+        i += 1
+    else:
+        print("Yeah not likely")
+    print(i)
+
+
+if __name__ == '__main__':
+    main()
+```
+#### Go
+```go
+package main
+
+import "fmt"
+
+func main() {
+	var x int
+	if func() bool {
+		for x = 0; x < 4; x++ {
+			if x == 5 {
+				return false
+			}
+		}
+		return true
+	}() {
+		fmt.Println("Well of course that didn't happen")
+	}
+	if func() bool {
+		for x = 0; x < 7; x++ {
+			if x == 5 {
+				return false
+			}
+		}
+		return true
+	}() {
+		fmt.Println("H-hey wait!")
+	}
+	i := 0
+	if func() bool {
+		for i < 3 {
+			fmt.Println("Works with while too")
+			for x = 0; x < 3; x++ {
+				fmt.Println("BTW don't worry about nested breaks")
+				break
+			}
+			if i == 10 {
+				return false
+			}
+			i += 1
+		}
+		return true
+	}() {
+		fmt.Println("Yeah not likely")
+	}
+	fmt.Println(i)
+}
+```
 ### algomajorityelement
 #### Python
 ```python
@@ -4442,76 +4593,6 @@ func main() {
 	}())
 }
 ```
-### scope
-#### Python
-```python
-import random
-
-
-def main():
-    if random.random() > 0.5:
-        a = 1
-    else:
-        a = 2
-
-    if random.random() > 0.5:
-        if random.random() > 0.5:
-            b = 1
-        else:
-            b = 2
-    else:
-        b = 3
-
-    def hello_world():
-        c = 3
-        print(c)
-
-    hello_world()
-    print(a, b)
-
-
-if __name__ == '__main__':
-    main()
-```
-#### Go
-```go
-package main
-
-import (
-	"fmt"
-	"math/rand"
-	"time"
-)
-
-func init() {
-	rand.Seed(time.Now().UnixNano())
-}
-
-func main() {
-	var a int
-	var b int
-	if rand.Float64() > 0.5 {
-		a = 1
-	} else {
-		a = 2
-	}
-	if rand.Float64() > 0.5 {
-		if rand.Float64() > 0.5 {
-			b = 1
-		} else {
-			b = 2
-		}
-	} else {
-		b = 3
-	}
-	hello_world := func() {
-		c := 3
-		fmt.Println(c)
-	}
-	hello_world()
-	fmt.Println(a, b)
-}
-```
 
 ## TODOs
 

diff --git a/examples/forelse.go b/examples/forelse.go
@@ -3,24 +3,43 @@ package main
 import "fmt"
 
 func main() {
-	broke := false
-	for x := 0; x < 4; x++ {
-		if x == 5 {
-			broke = true
-			break
+	var x int
+	if func() bool {
+		for x = 0; x < 4; x++ {
+			if x == 5 {
+				return false
+			}
 		}
-	}
-	if !broke {
+		return true
+	}() {
 		fmt.Println("Well of course that didn't happen")
 	}
-	broke = false
-	for x := 0; x < 7; x++ {
-		if x == 5 {
-			broke = true
-			break
+	if func() bool {
+		for x = 0; x < 7; x++ {
+			if x == 5 {
+				return false
+			}
 		}
-	}
-	if !broke {
+		return true
+	}() {
 		fmt.Println("H-hey wait!")
 	}
+	i := 0
+	if func() bool {
+		for i < 3 {
+			fmt.Println("Works with while too")
+			for x = 0; x < 3; x++ {
+				fmt.Println("BTW don't worry about nested breaks")
+				break
+			}
+			if i == 10 {
+				return false
+			}
+			i += 1
+		}
+		return true
+	}() {
+		fmt.Println("Yeah not likely")
+	}
+	fmt.Println(i)
 }
diff --git a/examples/forelse.py b/examples/forelse.py
@@ -11,6 +11,18 @@ def main():
     else:
         print("H-hey wait!")
 
+    i = 0
+    while i < 3:
+        print("Works with while too")
+        for x in range(3):
+            print("BTW don't worry about nested breaks")
+            break
+        if i == 10:
+            break
+        i += 1
+    else:
+        print("Yeah not likely")
+    print(i)
 
 
 if __name__ == '__main__':

diff --git a/pytago/go_ast/core.py b/pytago/go_ast/core.py
@@ -436,8 +436,9 @@ class GoAST(ast.AST):
 
     _prefix = "ast."
 
-    def __init__(self, parents=None, **kwargs):
+    def __init__(self, parents=None, _py_context=None, **kwargs):
         super().__init__(**kwargs)
+        self._py_context = _py_context or {}
         self.parents = parents or []
         for field_name in self._fields:
             field = getattr(self, field_name, None)
@@ -518,9 +519,8 @@ def __repr__(self):
 
 
 class Expr(GoAST):
-    def __init__(self, *args, _type_help=None, _py_context=None, **kwargs):
+    def __init__(self, *args, _type_help=None, **kwargs):
         self._type_help = _type_help
-        self._py_context = _py_context or {}
         super().__init__(**kwargs)
 
     def __or__(self, Y: 'Expr') -> 'BinaryExpr':
@@ -1226,14 +1226,12 @@ def __init__(self,
                  Kind: ObjKind = None,
                  Name: str = None,
                  Type: Expr = None,
-                 _py_context: dict = None,
                  **kwargs) -> None:
         self.Data = Data
         self.Decl = Decl
         self.Kind = Kind
         self.Name = Name
         self.Type = Type
-        self._py_context = _py_context or {}
         super().__init__(**kwargs)
 
 
@@ -2442,15 +2440,22 @@ def from_While(cls, node: ast.While):
         cond = build_expr_list([node.test])[0]
         match cond:
             case Ident(Name="true"):
-                return cls(Body=body)
+                loop = cls(Body=body)
             case BasicLit(Value=x) if not json.loads(x):
-                return cls(Body=body, Cond=Ident.from_str("false"))
+                loop = cls(Body=body, Cond=Ident.from_str("false"))
             case BasicLit(Kind=x) if x in [token.INT, token.STRING, token.FLOAT]:
-                return cls(Body=body)
+                loop = cls(Body=body)
             case Ident(Name="false") | Ident(Name="nil"):
-                return cls(Body=body, Cond=Ident.from_str("false"))
-
-        return cls(Body=body, Cond=cond)
+                loop = cls(Body=body, Cond=Ident.from_str("false"))
+            case _:
+                loop = cls(Body=body, Cond=cond)
+        if node.orelse:
+            loop._py_context["forelse"] = True
+            # The breaks of this loop can be replaced with a return False by a transformer
+            return IfStmt(Body=BlockStmt(List=build_stmt_list(node.orelse)),
+                          Cond=FuncLit(Body=BlockStmt(List=[loop, Ident("true").return_()]),
+                                       Type=FuncType(Results=FieldList(List=[Field(Type=GoBasicType.BOOL.ident)]))).call())
+        return loop
 
 
 class FuncType(Expr):
@@ -3107,6 +3112,13 @@ def from_For(cls, node: ast.For, **kwargs):
         key = Ident.from_str("_")
         value = build_expr_list([node.target])[0]
         x = build_expr_list([node.iter])[0]
+        if node.orelse:
+            for_body = cls(Body=body, Key=key, Tok=tok, Value=value, X=x, _py_context={**kwargs.get("_py_context", {}), "forelse": True},
+                           **kwargs)
+            # The breaks of this loop can be replaced with a return False by a transformer
+            return IfStmt(Body=BlockStmt(List=build_stmt_list(node.orelse)),
+                          Cond=FuncLit(Body=BlockStmt(List=[for_body, Ident("true").return_()]),
+                                       Type=FuncType(Results=FieldList(List=[Field(Type=GoBasicType.BOOL.ident)]))).call())
         return cls(Body=body, Key=key, Tok=tok, Value=value, X=x, **kwargs)
 
     @classmethod