From 500605badcdbc571c8b6f85043aeef25d5fadf49 Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Tue, 17 May 2022 13:57:38 +0200
Subject: [PATCH] Add an action to update the method signature in a docstring
 attached to the method.

---
 src/requests/actions.jl  | 37 +++++++++++++++++++++++++++++++++++++
 test/requests/actions.jl | 28 ++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/src/requests/actions.jl b/src/requests/actions.jl
index 20da0322..81b6f329 100644
--- a/src/requests/actions.jl
+++ b/src/requests/actions.jl
@@ -615,6 +615,34 @@ function add_docstring_template(x, _, conn)
     return
 end
 
+function update_docstring_sig(x, _, conn)
+    is_in_function_signature(x, nothing; with_docstring=true) || return
+    func = _get_parent_fexpr(x, CSTParser.defines_function)
+    # Current docstring
+    docstr_expr = func.parent.args[3]
+    docstr = valof(docstr_expr)
+    file, docstr_offset = get_file_loc(docstr_expr)
+    # New signature in the code
+    sig = func.args[1]
+    _, sig_offset = get_file_loc(sig)
+    sig_str = get_text(file)[sig_offset .+ (1:sig.span)]
+    # Heuristic for finding a signature in the current docstring
+    reg = r"\A    .*$"m
+    if (m = match(reg, valof(docstr_expr)); m !== nothing)
+        docstr = replace(docstr, reg => string("    ", sig_str))
+    else
+        docstr = string("    ", sig_str, "\n\n", docstr)
+    end
+    newline = endswith(docstr, "\n") ? "" : "\n"
+    # Rewrap in """"
+    docstr = string("\"\"\"\n", docstr, newline, "\"\"\"")
+    tde = TextDocumentEdit(VersionedTextDocumentIdentifier(get_uri(file), get_version(file)), TextEdit[
+        TextEdit(Range(file, docstr_offset .+ (0:docstr_expr.span)), docstr)
+    ])
+    JSONRPC.send(conn, workspace_applyEdit_request_type, ApplyWorkspaceEditParams(missing, WorkspaceEdit(missing, TextDocumentEdit[tde])))
+    return
+end
+
 # Adding a CodeAction requires defining:
 # * a unique id
 # * a description
@@ -730,3 +758,12 @@ LSActions["AddDocstringTemplate"] = ServerAction(
     is_in_function_signature,
     add_docstring_template,
 )
+
+LSActions["UpdateDocstringSignature"] = ServerAction(
+    "UpdateDocstringSignature",
+    "Update method signature in docstring",
+    missing,
+    missing,
+    (args...) -> is_in_function_signature(args...; with_docstring=true),
+    update_docstring_sig,
+)
diff --git a/test/requests/actions.jl b/test/requests/actions.jl
index c2760c5c..6dbbcc00 100644
--- a/test/requests/actions.jl
+++ b/test/requests/actions.jl
@@ -193,3 +193,31 @@ end
     c = filter(c -> c.command == "AddDocstringTemplate", action_request_test(0, 1))[1]
     LanguageServer.workspace_executeCommand_request(LanguageServer.ExecuteCommandParams(missing, c.command, c.arguments), server, server.jr_endpoint)
 end
+
+@testset "Update docstring signature" begin
+    doc = settestdoc("""
+        "hello"
+        f(x) = x
+
+        \"\"\"hello\"\"\"
+        g(x) = x
+
+        \"\"\"
+            h()
+
+        hello
+        \"\"\"
+        function h(x)
+        end
+
+        i(x) = x
+        """)
+
+    @test any(c.command == "UpdateDocstringSignature" for c in action_request_test(1, 0))
+    @test any(c.command == "UpdateDocstringSignature" for c in action_request_test(4, 0))
+    @test any(c.command == "UpdateDocstringSignature" for c in action_request_test(11, 0))
+    @test !any(c.command == "UpdateDocstringSignature" for c in action_request_test(14, 0))
+
+    c = filter(c -> c.command == "UpdateDocstringSignature", action_request_test(1, 0))[1]
+    LanguageServer.workspace_executeCommand_request(LanguageServer.ExecuteCommandParams(missing, c.command, c.arguments), server, server.jr_endpoint)
+end