Fix indent of array comprehensions, fixes #20

src/runestone.jl

@@ -2295,7 +2295,7 @@ end
 
 # Literal array nodes and also ref-nodes (which can be either a typed-array or a getindex)
 function indent_array(ctx::Context, node::Node)
-    @assert kind(node) in KSet"vect vcat typed_vcat ncat ref"
+    @assert kind(node) in KSet"vect vcat typed_vcat ncat ref comprehension typed_comprehension"
     kids = verified_kids(node)
     opening_bracket_idx = findfirst(x -> kind(x) === K"[", kids)::Int
     closing_bracket_idx = findnext(x -> kind(x) === K"]", kids, opening_bracket_idx + 1)::Int
@@ -2429,7 +2429,7 @@ function insert_delete_mark_newlines(ctx::Context, node::Node)
         return indent_do(ctx, node)
     elseif is_paren_block(node)
         return indent_paren_block(ctx, node)
-    elseif kind(node) in KSet"vect vcat typed_vcat ncat ref"
+    elseif kind(node) in KSet"vect vcat typed_vcat ncat ref comprehension typed_comprehension"
         return indent_array(ctx, node)
     elseif kind(node) in KSet"row"
         return indent_array_row(ctx, node)

test/runtests.jl

@@ -670,6 +670,13 @@ end
             @test format_string("$(t)[a b\n$(sp)c d]") == "$(t)[\n    a b\n    c d\n]"
             @test format_string("$(t)[\n$(sp)a b\n$(sp)c d\n$(sp)]") == "$(t)[\n    a b\n    c d\n]"
         end
+        # array comprehension
+        for t in ("", "T")
+            @test format_string("$(t)[$(sp)a for a in b$(sp)\n$(sp)]") ==
+                format_string("$(t)[$(sp)\n$(sp)a for a in b$(sp)]") ==
+                format_string("$(t)[$(sp)\n$(sp)a for a in b$(sp)\n$(sp)]") ==
+                "$(t)[\n    a for a in b\n]"
+        end
         # Single line begin-end
         @test format_string("begin x\n$(sp)end") == "begin x\nend"
         @test format_string("begin x end") == "begin x end"