resolution for 741 - indent function and macro definition arguments b…

…y 2x the indent
domluna · Oct 28, 2023 · d4e4c4c · d4e4c4c
1 parent ce4ad47
commit d4e4c4c
Show file tree

Hide file tree

Showing 2 changed files with 110 additions and 2 deletions.
diff --git a/src/styles/sciml/pretty.jl b/src/styles/sciml/pretty.jl
@@ -207,3 +207,49 @@ function p_unaryopcall(ds::SciMLStyle, cst::CSTParser.EXPR, s::State; kwargs...)
     end
     t
 end
+
+function p_functiondef(ds::SciMLStyle, cst::CSTParser.EXPR, s::State)
+    style = getstyle(ds)
+    t = FST(FunctionN, cst, nspaces(s))
+    add_node!(t, pretty(style, cst[1], s), s)
+    add_node!(t, Whitespace(1), s)
+    s.indent += s.opts.indent
+    add_node!(t, pretty(style, cst[2], s), s, join_lines = true)
+    s.indent -= s.opts.indent
+    if length(cst) > 3
+        if cst[3].fullspan == 0
+            n = pretty(style, cst[4], s)
+            if s.opts.join_lines_based_on_source
+                join_lines = t.endline == n.startline
+                join_lines && (add_node!(t, Whitespace(1), s))
+                add_node!(t, n, s, join_lines = join_lines)
+            else
+                add_node!(t, Whitespace(1), s)
+                add_node!(t, n, s, join_lines = true)
+            end
+        else
+            s.indent += s.opts.indent
+            n = pretty(style, cst[3], s, ignore_single_line = true)
+            s.opts.always_use_return && prepend_return!(n, s)
+            add_node!(t, n, s, max_padding = s.opts.indent)
+            s.indent -= s.opts.indent
+            add_node!(t, pretty(style, cst[4], s), s)
+        end
+    else
+        # function stub `function foo end`
+        n = pretty(style, cst[3], s)
+        if s.opts.join_lines_based_on_source
+            join_lines = t.endline == n.startline
+            join_lines && (add_node!(t, Whitespace(1), s))
+            add_node!(t, n, s, join_lines = join_lines)
+        else
+            add_node!(t, Whitespace(1), s)
+            add_node!(t, n, s, join_lines = true)
+        end
+    end
+    t
+end
+
+function p_macro(ds::SciMLStyle, cst::CSTParser.EXPR, s::State)
+    p_functiondef(ds, cst, s)
+end
diff --git a/test/sciml_style.jl b/test/sciml_style.jl
@@ -35,8 +35,8 @@
     """
     formatted_str = raw"""
     function BipartiteGraph(fadj::AbstractVector,
-        badj::Union{AbstractVector, Integer} = maximum(maximum, fadj);
-        metadata = nothing)
+            badj::Union{AbstractVector, Integer} = maximum(maximum, fadj);
+            metadata = nothing)
         BipartiteGraph(mapreduce(length, +, fadj; init = 0), fadj, badj, metadata)
     end
     """
@@ -293,4 +293,66 @@
     @test format_text(str, SciMLStyle(), variable_call_indent = ["test"]) == formatted_str
     @test format_text(str, SciMLStyle(), variable_call_indent = ["SVector", "test"]) ==
           formatted_str
+
+    str = """
+    function alg_cache(alg::FineRKN4, u, rate_prototype, ::Type{uEltypeNoUnits},
+           ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
+          dt, reltol, p, calck, ::Val{true}) where {uEltypeNoUnits, 
+        uBottomEltypeNoUnits,tTypeNoUnits}
+
+        reduced_rate_prototype = rate_prototype.x[2]
+        tab = FineRKN4ConstantCache(constvalue(uBottomEltypeNoUnits), constvalue(tTypeNoUnits))
+        k1 = zero(rate_prototype)
+        k2 = zero(reduced_rate_prototype)
+        k3 = zero(reduced_rate_prototype)
+        k4 = zero(reduced_rate_prototype)
+        k5 = zero(reduced_rate_prototype)
+        k = zero(rate_prototype)
+        utilde = zero(u)
+        atmp = similar(u, uEltypeNoUnits)
+        recursivefill!(atmp, false)
+        tmp = zero(u)
+        FineRKN4Cache(u, uprev, k1, k2, k3, k4, k5, k, utilde, tmp, atmp, tab)
+    end"""
+
+    formatted_str = """
+    function alg_cache(alg::FineRKN4, u, rate_prototype, ::Type{uEltypeNoUnits},
+            ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
+            dt, reltol, p, calck,
+            ::Val{true}) where {uEltypeNoUnits,
+            uBottomEltypeNoUnits, tTypeNoUnits}
+        reduced_rate_prototype = rate_prototype.x[2]
+        tab = FineRKN4ConstantCache(constvalue(uBottomEltypeNoUnits), constvalue(tTypeNoUnits))
+        k1 = zero(rate_prototype)
+        k2 = zero(reduced_rate_prototype)
+        k3 = zero(reduced_rate_prototype)
+        k4 = zero(reduced_rate_prototype)
+        k5 = zero(reduced_rate_prototype)
+        k = zero(rate_prototype)
+        utilde = zero(u)
+        atmp = similar(u, uEltypeNoUnits)
+        recursivefill!(atmp, false)
+        tmp = zero(u)
+        FineRKN4Cache(u, uprev, k1, k2, k3, k4, k5, k, utilde, tmp, atmp, tab)
+    end"""
+    @test format_text(str, SciMLStyle()) == formatted_str
+
+    str = """
+    function SpatialMassActionJump(urates::A, srates::B, rs::S, ns::U, pmapper::V;
+                                   scale_rates = true, useiszero = true,
+                                   nocopy = false) where {A <: AVecOrNothing,
+                                                          B <: AMatOrNothing, S, U, V}
+        SpatialMassActionJump{A, B, S, U, V}(urates, srates, rs, ns, pmapper, scale_rates,
+                                             useiszero, nocopy)
+    end"""
+
+    formatted_str = """
+    function SpatialMassActionJump(urates::A, srates::B, rs::S, ns::U, pmapper::V;
+            scale_rates = true, useiszero = true,
+            nocopy = false) where {A <: AVecOrNothing,
+            B <: AMatOrNothing, S, U, V}
+        SpatialMassActionJump{A, B, S, U, V}(urates, srates, rs, ns, pmapper, scale_rates,
+            useiszero, nocopy)
+    end"""
+    @test format_text(str, SciMLStyle()) == formatted_str
 end