diff --git a/src/runestone.jl b/src/runestone.jl index 3a12fe8..0e2d2ff 100644 --- a/src/runestone.jl +++ b/src/runestone.jl @@ -457,23 +457,26 @@ end # replace `=` and `∈` with `in` in for-loops function for_loop_use_in(ctx::Context, node::JuliaSyntax.GreenNode) - if !(JuliaSyntax.kind(node) === K"for" && !is_leaf(node) && n_children(node) == 4) + if !( + (JuliaSyntax.kind(node) === K"for" && !is_leaf(node) && n_children(node) == 4) || + (JuliaSyntax.kind(node) === K"generator" && n_children(node) == 3) # TODO: Unsure about 3. + ) return nothing end pos = position(ctx.fmt_io) # In case a reset is needed later bytes = node_bytes(ctx, node) children = verified_children(node) - for_index = 1 + for_index = findfirst(c -> JuliaSyntax.kind(c) === K"for" && is_leaf(c), children)::Int for_node = children[for_index] - # TODO: Could there be whitespace here before the K"for"? @assert JuliaSyntax.kind(for_node) === K"for" && JuliaSyntax.span(for_node) == 3 && is_leaf(for_node) && JuliaSyntax.is_trivia(for_node) - accept_node!(ctx, for_node) + for i in 1:for_index + accept_node!(ctx, children[i]) + end # The for loop specification node can be either K"=" or K"cartesian_iterator" - for_spec_index = 2 + for_spec_index = for_index + 1 for_spec_node = children[for_spec_index] @assert JuliaSyntax.kind(for_spec_node) in KSet"= cartesian_iterator" - pos_before = position(ctx.fmt_io) if JuliaSyntax.kind(for_spec_node) === K"=" for_spec_node′ = replace_with_in(ctx, for_spec_node) else @@ -485,7 +488,7 @@ function for_loop_use_in(ctx::Context, node::JuliaSyntax.GreenNode) seek(ctx.fmt_io, pos) return nothing end - @assert position(ctx.fmt_io) == pos + JuliaSyntax.span(for_node) + JuliaSyntax.span(for_spec_node′) + @assert position(ctx.fmt_io) == pos + mapreduce(JuliaSyntax.span, +, @view(children[1:for_index])) + JuliaSyntax.span(for_spec_node′) # Insert the new for spec node children′ = copy(children) children′[for_spec_index] = for_spec_node′ diff --git a/test/runtests.jl b/test/runtests.jl index efed991..6cdf3d5 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -226,11 +226,22 @@ end @test format_string("a >: T >: S") == "a >: T >: S" end -@testset "replace ∈ and = with in in for loops" begin +@testset "replace ∈ and = with in in for loops and generators" begin for sp in ("", " ", " "), op in ("∈", "=", "in") op == "in" && sp == "" && continue + # for loops @test format_string("for i$(sp)$(op)$(sp)I\nend") == "for i in I\nend" @test format_string("for i$(sp)$(op)$(sp)I, j$(sp)$(op)$(sp)J\nend") == "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 (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)") == + "$(l)(i, j) for i in I, j in J$(r)" + @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 end end