Repository: domluna/JuliaFormatter.jl (julia)

Issue:
YASStyle doesn't work with for comprehensions & `=`
```
src = "[x[i] for i = 1:length(x)]"

format_text(src, YASStyle())
ERROR: KeyError: key 27:27 not found
```
Using `in` works!

Pull request description:
fixes #320

Patch:
diff --git a/src/styles/default/pretty.jl b/src/styles/default/pretty.jl
index 38307ed59..29b1d03cc 100644
--- a/src/styles/default/pretty.jl
+++ b/src/styles/default/pretty.jl
@@ -1005,28 +1005,6 @@ p_let(style::S, cst::CSTParser.EXPR, s::State) where {S<:AbstractStyle} =
 # for i = 1:10 body end
 #
 # https://github.com/domluna/JuliaFormatter.jl/issues/34
-function eq_to_in_normalization!(cst::CSTParser.EXPR, always_for_in::Bool)
-    if cst.typ === CSTParser.BinaryOpCall
-        op = cst[2]
-        rhs = cst[3]
-
-        if always_for_in
-            cst[2].kind = Tokens.IN
-            return
-        end
-
-        if op.kind === Tokens.EQ && !is_colon_op(rhs)
-            cst[2].kind = Tokens.IN
-        elseif op.kind === Tokens.IN && is_colon_op(rhs)
-            cst[2].kind = Tokens.EQ
-        end
-    elseif cst.typ === CSTParser.Block || cst.typ === CSTParser.InvisBrackets
-        for a in cst
-            eq_to_in_normalization!(a, always_for_in)
-        end
-    end
-end
-
 function eq_to_in_normalization!(fst::FST, always_for_in::Bool)
     if fst.typ === CSTParser.BinaryOpCall
         idx = findfirst(n -> n.typ === CSTParser.OPERATOR, fst.nodes)
@@ -1059,9 +1037,6 @@ function p_for(ds::DefaultStyle, cst::CSTParser.EXPR, s::State)
     t = FST(cst, nspaces(s))
     add_node!(t, pretty(style, cst[1], s), s)
     add_node!(t, Whitespace(1), s)
-    # if cst[1].kind === Tokens.FOR
-    #     eq_to_in_normalization!(cst[2], s.opts.always_for_in)
-    # end
 
     n = if cst[2].typ === CSTParser.Block
         s.indent += s.opts.indent
@@ -2011,6 +1986,7 @@ function p_generator(ds::DefaultStyle, cst::CSTParser.EXPR, s::State)
         else
             add_node!(t, n, s, join_lines = true)
         end
+
         has_for_kw && eq_to_in_normalization!(n, s.opts.always_for_in)
     end
     t
diff --git a/src/styles/yas/pretty.jl b/src/styles/yas/pretty.jl
index 7de459499..b49c8ff06 100644
--- a/src/styles/yas/pretty.jl
+++ b/src/styles/yas/pretty.jl
@@ -291,10 +291,15 @@ end
 function p_generator(ys::YASStyle, cst::CSTParser.EXPR, s::State)
     style = getstyle(ys)
     t = FST(cst, nspaces(s))
+    has_for_kw = false
 
     for (i, a) in enumerate(cst)
         n = pretty(style, a, s)
         if a.typ === CSTParser.KEYWORD
+            if !has_for_kw && a.kind === Tokens.FOR
+                has_for_kw = true
+            end
+
             incomp = parent_is(
                 a,
                 is_iterable,
@@ -311,16 +316,15 @@ function p_generator(ys::YASStyle, cst::CSTParser.EXPR, s::State)
             add_node!(t, n, s, join_lines = true)
             add_node!(t, Whitespace(1), s)
 
-            if a.kind === Tokens.FOR
-                for j = i+1:length(cst)
-                    eq_to_in_normalization!(cst[j], s.opts.always_for_in)
-                end
-            end
-
             if !is_gen(cst.args[i+1])
                 tup = p_tupleh(style, cst.args[i+1:length(cst)], s)
                 add_node!(t, tup, s, join_lines = true)
-                return t
+                if has_for_kw
+                    for nn in tup.nodes
+                        eq_to_in_normalization!(nn, s.opts.always_for_in)
+                    end
+                end
+                break
             end
         elseif CSTParser.is_comma(a) && i < length(cst) && !is_punc(cst[i+1])
             add_node!(t, n, s, join_lines = true)
@@ -328,6 +332,8 @@ function p_generator(ys::YASStyle, cst::CSTParser.EXPR, s::State)
         else
             add_node!(t, n, s, join_lines = true)
         end
+
+        has_for_kw && eq_to_in_normalization!(n, s.opts.always_for_in)
     end
 
     t


Test patch:
diff --git a/test/yas_style.jl b/test/yas_style.jl
index 66bb32121..0df422a11 100644
--- a/test/yas_style.jl
+++ b/test/yas_style.jl
@@ -435,4 +435,10 @@
         end"""
         @test yasfmt(str_, 4, 92) == str
     end
+
+    @testset "issue 320" begin
+        str_ = "[x[i] for i = 1:length(x)]"
+        str = "[x[i] for i in 1:length(x)]"
+        @test yasfmt(str_, 4, 92, always_for_in = true) == str
+    end
 end