fix generators/flatten #41

src/components/loops.jl

@@ -70,7 +70,7 @@ function parse_generator(ps::ParseState, ret)
         end
         next(ps)
         unshift!(ranges, INSTANCE(ps))
-        @catcherror ps cond = @closer ps paren parse_expression(ps)
+        @catcherror ps cond = @closer ps range @closer ps paren parse_expression(ps)
         unshift!(ranges, cond)
         push!(ret, ranges)
     else

src/conversion.jl

@@ -619,27 +619,48 @@ function Expr(x::EXPR{Comprehension})
     end
     ret
 end
-# NEEDS FIX (only works for 2 levels )
+
 function Expr(x::EXPR{Flatten})
-    ret = Expr(:flatten, Expr(x.args[1]))
+    iters, args = get_inner_gen(x)
+    i = shift!(iters)
+    ex = Expr(:generator, Expr(args[1]), convert_iter_assign(i))
+    for i in iters
+        ex = Expr(:generator, ex, convert_iter_assign(i))
+        ex = Expr(:flatten, ex)
+    end
+    # ret = Expr(:flatten, ex)
 
-    gen1 = ret.args[1].args[1].args[2]
-    gen2 = ret.args[1].args[2]
-    ret.args[1].args[1].args[2] = gen2
-    ret.args[1].args[2] = gen1
-    return ret
+    return ex
+end
+
+
+function get_inner_gen(x, iters = [], arg = []) iters, arg end
+function get_inner_gen(x::EXPR{Flatten}, iters = [], arg = [])
+    get_inner_gen(x.args[1], iters, arg)
+    iters, arg
+end
+function get_inner_gen(x::EXPR{Generator}, iters = [], arg = [])
+    push!(iters, get_iter(x))
+    if x.args[1] isa EXPR{Generator} || x.args[1] isa EXPR{Flatten}
+        get_inner_gen(x.args[1], iters, arg)
+    else
+        push!(arg, x.args[1])
+    end
+    iters, arg
+end
+
+
+function get_iter(x) end
+function get_iter(x::EXPR{Generator})
+    return x.args[3]
 end
 
 function Expr(x::EXPR{Generator})
     ret = Expr(:generator, Expr(x.args[1]))
     for i = 3:length(x.args)
         a = x.args[i]
         if !(a isa EXPR{P} where P <: PUNCTUATION)
-            if a isa EXPR{BinaryOpCall} && (a.args[2] isa EXPR{OPERATOR{ComparisonOp,Tokens.IN,false}} || a.args[2] isa EXPR{OPERATOR{ComparisonOp,Tokens.ELEMENT_OF,false}})
-                push!(ret.args, Expr(:(=), Expr(a.args[1]), Expr(a.args[3])))
-            else
-                push!(ret.args, Expr(a))
-            end
+            push!(ret.args, convert_iter_assign(a))
         end
     end
     ret
@@ -648,18 +669,20 @@ end
 function Expr(x::EXPR{Filter})
     ret = Expr(:filter)
     for a in x.args
-        if !(a isa EXPR{KEYWORD{Tokens.IF}})
-            if a isa EXPR{BinaryOpCall} && (a.args[2] isa EXPR{OPERATOR{ComparisonOp,Tokens.IN,false}} || a.args[2] isa EXPR{OPERATOR{ComparisonOp,Tokens.ELEMENT_OF,false}})
-                push!(ret.args, Expr(:(=), Expr(a.args[1]), Expr(a.args[3])))
-            else
-                push!(ret.args, Expr(a))
-            end
+        if !(a isa EXPR{KEYWORD{Tokens.IF}} || a isa EXPR{<:PUNCTUATION})
+            push!(ret.args, convert_iter_assign(a))
         end
     end
     ret
 end
 
-
+function convert_iter_assign(a)
+    if a isa EXPR{BinaryOpCall} && (a.args[2] isa EXPR{OPERATOR{ComparisonOp,Tokens.IN,false}} || a.args[2] isa EXPR{OPERATOR{ComparisonOp,Tokens.ELEMENT_OF,false}})
+        return Expr(:(=), Expr(a.args[1]), Expr(a.args[3]))
+    else
+        return Expr(a)
+    end
+end
 
 
 function Expr(x::EXPR{TypedComprehension})

src/utils.jl

@@ -23,7 +23,7 @@ function closer(ps::ParseState)
     (ps.closer.ifelse && ps.nt.kind == Tokens.ELSEIF || ps.nt.kind == Tokens.ELSE) ||
     (ps.closer.ifop && isoperator(ps.nt) && (precedence(ps.nt) <= 0 || ps.nt.kind == Tokens.COLON)) ||
     (ps.closer.trycatch && (ps.nt.kind == Tokens.CATCH || ps.nt.kind == Tokens.FINALLY || ps.nt.kind == Tokens.END)) ||
-    (ps.closer.range && (ps.nt.kind == Tokens.FOR || iscomma(ps.nt))) ||
+    (ps.closer.range && (ps.nt.kind == Tokens.FOR || iscomma(ps.nt) || ps.nt.kind == Tokens.IF)) ||
     (ps.closer.ws && !isemptyws(ps.ws) &&
         !(ps.nt.kind == Tokens.COMMA) &&
         !(ps.t.kind == Tokens.COMMA) &&

test/parser.jl

@@ -331,6 +331,17 @@ end
     """ |> test_expr
     @test "[ V[j][i]::T for i=1:length(V[1]), j=1:length(V) ]" |> test_expr
     @test "all(d ≥ 0 for d in B.dims)" |> test_expr
+    @test "(arg for x in X)" |> test_expr
+    @test "(arg for x in X for y in Y)" |> test_expr
+    @test "(arg for x in X for y in Y for z in Z)" |> test_expr
+    @test "(arg for x in X if A)" |> test_expr
+    @test "(arg for x in X if A for y in Y)" |> test_expr
+    @test "(arg for x in X if A for y in Y if B)" |> test_expr
+    @test "(arg for x in X if A for y in Y for z in Z)" |> test_expr
+    @test "(arg for x in X if A for y in Y if B for z in Z)" |> test_expr
+    @test "(arg for x in X if A for y in Y if B for z in Z if C)" |> test_expr
+    @test_broken "(arg for x in X, y in Y for z in Z)" |> test_expr
+    @test "(arg for x in X, y in Y if A for z in Z)" |> test_expr
 end
 
 @testset "Macros " begin