Fix replacement of = and \in in filtered generators

src/runestone.jl

@@ -945,6 +945,7 @@ function spaces_around_keywords(ctx::Context, node::Node)
 end
 
 # Replace the K"=" operator with `in`
+# TODO: This method doesn't reset the stream so callers should not accept_node??
 function replace_with_in(ctx::Context, node::Node)
     @assert kind(node) === K"=" && !is_leaf(node) && meta_nargs(node) == 3
     kids = verified_kids(node)
@@ -982,6 +983,28 @@ function replace_with_in(ctx::Context, node::Node)
     return make_node(node, kids′)
 end
 
+function replace_with_in_filter(ctx::Context, node::Node)
+    @assert kind(node) === K"filter" && !is_leaf(node)
+    pos = position(ctx.fmt_io)
+    kids = verified_kids(node)
+    idx = findfirst(x -> kind(x) === K"=" && !is_leaf(x), kids)::Int
+    for i in 1:(idx - 1)
+        accept_node!(ctx, kids[i])
+    end
+    kid = kids[idx]
+    kid′ = replace_with_in(ctx, kid)
+    if kid′ === nothing
+        seek(ctx.fmt_io, pos)
+        return nothing
+    end
+    kids = copy(kids)
+    kids[idx] = kid′
+    for i in (idx + 1):length(kids)
+        accept_node!(ctx, kids[i])
+    end
+    return make_node(node, kids)
+end
+
 function replace_with_in_cartesian(ctx::Context, node::Node)
     @assert kind(node) === K"cartesian_iterator" && !is_leaf(node)
     kids = verified_kids(node)
@@ -1029,9 +1052,11 @@ function for_loop_use_in(ctx::Context, node::Node)
     # The for loop specification node can be either K"=" or K"cartesian_iterator"
     for_spec_index = for_index + 1
     for_spec_node = kids[for_spec_index]
-    @assert kind(for_spec_node) in KSet"= cartesian_iterator"
+    @assert kind(for_spec_node) in KSet"= cartesian_iterator filter"
     if kind(for_spec_node) === K"="
         for_spec_node′ = replace_with_in(ctx, for_spec_node)
+    elseif kind(for_spec_node) === K"filter"
+        for_spec_node′ = replace_with_in_filter(ctx, for_spec_node)
     else
         @assert kind(for_spec_node) === K"cartesian_iterator"
         for_spec_node′ = replace_with_in_cartesian(ctx, for_spec_node)

test/runtests.jl

@@ -389,7 +389,7 @@ end
             "for i in I, j in J\nend"
         @test format_string("for i$(sp)$(op)$(sp)I, j$(sp)$(op)$(sp)J, k$(sp)$(op)$(sp)K\nend") ==
             "for i in I, j in J, k in K\nend"
-        # for generators
+        # for generators, filter
         for (l, r) in (("[", "]"), ("(", ")"))
             @test format_string("$(l)i for i$(sp)$(op)$(sp)I$(r)") == "$(l)i for i in I$(r)"
             @test format_string("$(l)(i, j) for i$(sp)$(op)$(sp)I, j$(sp)$(op)$(sp)J$(r)") ==
@@ -397,6 +397,11 @@ end
             @test format_string("$(l)(i, j, k) for i$(sp)$(op)$(sp)I, j$(sp)$(op)$(sp)J, k$(sp)$(op)$(sp)K$(r)") ==
                 "$(l)(i, j, k) for i in I, j in J, k in K$(r)"
         end
+        # K"filter"
+        for (l, r) in (("[", "]"), ("(", ")"))
+            @test format_string("$(l)i for i$(sp)$(op)$(sp)I if i < 2$(r)") ==
+                "$(l)i for i in I if i < 2$(r)"
+        end
     end
 end