Repository: domluna/JuliaFormatter.jl
Pull request #791 (2023-12-12 07:39:38)

SciML function indent


Patch:
diff --git a/Project.toml b/Project.toml
index 2f4eb596..25f00f1a 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "JuliaFormatter"
 uuid = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 authors = ["Dominique Luna <dluna132@gmail.com>"]
-version = "1.0.44"
+version = "1.0.45"
 
 [deps]
 CSTParser = "00ebfdb7-1f24-5e51-bd34-a7502290713f"
diff --git a/src/nest_utils.jl b/src/nest_utils.jl
index d4fb8ab9..01773b1b 100644
--- a/src/nest_utils.jl
+++ b/src/nest_utils.jl
@@ -47,11 +47,11 @@ end
 
 function add_indent!(fst::FST, s::State, indent)
     indent == 0 && return
-    lo = s.line_offset
     f = (fst::FST, s::State) -> begin
         fst.indent += indent
         return nothing
     end
+    lo = s.line_offset
     walk(f, fst, s)
     s.line_offset = lo
 end
diff --git a/src/styles/default/nest.jl b/src/styles/default/nest.jl
index 85012da1..e71b04ef 100644
--- a/src/styles/default/nest.jl
+++ b/src/styles/default/nest.jl
@@ -145,6 +145,10 @@ function nest!(ds::DefaultStyle, fst::FST, s::State)
         n_unaryopcall!(style, fst, s)
     elseif fst.typ === StringN
         n_string!(style, fst, s)
+    elseif fst.typ === FunctionN
+        n_functiondef!(style, fst, s)
+    elseif fst.typ === Macro
+        n_macro!(style, fst, s)
     else
         nest!(style, fst.nodes::Vector, s, fst.indent, extra_margin = fst.extra_margin)
     end
@@ -152,6 +156,16 @@ end
 nest!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
     nest!(DefaultStyle(style), fst, s)
 
+function n_functiondef!(ds::DefaultStyle, fst::FST, s::State)
+    nest!(ds, fst.nodes::Vector, s, fst.indent, extra_margin = fst.extra_margin)
+end
+n_functiondef!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
+    n_functiondef!(DefaultStyle(style), fst, s)
+
+n_macro!(ds::DefaultStyle, fst::FST, s::State) = n_functiondef!(ds, fst, s)
+n_macro!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
+    n_macro!(DefaultStyle(style), fst, s)
+
 function n_string!(ds::DefaultStyle, fst::FST, s::State)
     style = getstyle(ds)
     # difference in positioning of the string
diff --git a/src/styles/sciml/nest.jl b/src/styles/sciml/nest.jl
index 9a1ad9fc..2f177fb8 100644
--- a/src/styles/sciml/nest.jl
+++ b/src/styles/sciml/nest.jl
@@ -52,3 +52,44 @@ function n_binaryopcall!(ss::SciMLStyle, fst::FST, s::State; indent::Int = -1)
     walk(increment_line_offset!, (fst.nodes::Vector)[1:end-1], s, fst.indent)
     nest!(style, fst[end], s)
 end
+
+function n_functiondef!(ss::SciMLStyle, fst::FST, s::State)
+    style = getstyle(ss)
+    if s.opts.yas_style_nesting
+        nest!(
+            YASStyle(style),
+            fst.nodes::Vector,
+            s,
+            fst.indent,
+            extra_margin = fst.extra_margin,
+        )
+    else
+        nest!(
+            DefaultStyle(style),
+            fst.nodes::Vector,
+            s,
+            fst.indent,
+            extra_margin = fst.extra_margin,
+        )
+
+        base_indent = fst.indent
+        closers = FST[]
+        f = (fst::FST, s::State) -> begin
+            if is_closer(fst) && fst.indent == base_indent
+                push!(closers, fst)
+            end
+            fst.indent += s.opts.indent
+            return nothing
+        end
+        lo = s.line_offset
+        walk(f, fst[3], s)
+        s.line_offset = lo
+        for c in closers
+            c.indent -= s.opts.indent
+        end
+    end
+end
+
+function n_macro!(ss::SciMLStyle, fst::FST, s::State)
+    n_functiondef!(ss, fst, s)
+end
diff --git a/src/styles/sciml/pretty.jl b/src/styles/sciml/pretty.jl
index 0e804571..beffdcab 100644
--- a/src/styles/sciml/pretty.jl
+++ b/src/styles/sciml/pretty.jl
@@ -208,52 +208,6 @@ function p_unaryopcall(ds::SciMLStyle, cst::CSTParser.EXPR, s::State; kwargs...)
     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
-
 function p_curly(ss::SciMLStyle, cst::CSTParser.EXPR, s::State)
     style = getstyle(ss)
     if s.opts.yas_style_nesting


Test patch:
diff --git a/test/sciml_style.jl b/test/sciml_style.jl
index 3d89595b..fbeb58ab 100644
--- a/test/sciml_style.jl
+++ b/test/sciml_style.jl
@@ -373,4 +373,63 @@
     #                               Foo
     #                              }
     @test format_text(str, SciMLStyle(), yas_style_nesting = true) == str
+
+    # https://github.com/SciML/JumpProcesses.jl/pull/362#discussion_r1375432275
+    str_ = """
+    function (p::CoevolveJumpAggregation{
+        T,
+        S,
+        F1,
+        F2,
+    })(integrator::AbstractSSAIntegrator) where {T, S, F1, F2 <: Union{Tuple, Nothing}}
+        body
+    end"""
+    str = """
+    function (p::CoevolveJumpAggregation{
+            T,
+            S,
+            F1,
+            F2,
+    })(integrator::AbstractSSAIntegrator) where {
+            T,
+            S,
+            F1,
+            F2 <:
+            Union{
+                Tuple,
+                Nothing,
+            },
+    }
+        body
+    end"""
+    @test format_text(str_, SciMLStyle(), margin = 1) == str
+
+    str_ = """
+    macro (p::CoevolveJumpAggregation{
+        T,
+        S,
+        F1,
+        F2,
+    })(integrator::AbstractSSAIntegrator) where {T, S, F1, F2 <: Union{Tuple, Nothing}}
+        body
+    end"""
+    str = """
+    macro (p::CoevolveJumpAggregation{
+            T,
+            S,
+            F1,
+            F2,
+    })(integrator::AbstractSSAIntegrator) where {
+            T,
+            S,
+            F1,
+            F2 <:
+            Union{
+                Tuple,
+                Nothing,
+            },
+    }
+        body
+    end"""
+    @test format_text(str_, SciMLStyle(), margin = 1) == str
 end