feat: add undojoin as a format option (#488)

* feat: add new format option: undojoin

This option allow user to automatically perform undojoin command before
formatting. This is useful if the user uses an autosave plugin + format
on save, because in case of undo, it will undo the last change AND the
formatting. Without this option, it will only undo the formatting.

* fix: passed linting

* fix: apply undojoin for LSP formatting

* doc: fix type annotations for apply_format

* doc: regenerate documentation

---------

Co-authored-by: Steven Arcangeli <stevearc@stevearc.com>

6 files changed, 50 insertions, 13 deletions

diff --git a/README.md b/README.md
index 8fab003..82f5436 100644
--- a/README.md
+++ b/README.md
@@ -593,6 +593,7 @@ Format a buffer
 |          | bufnr                                                | `nil\|integer`                       | Format this buffer (default 0)                                                                                                                       |
 |          | async                                                | `nil\|boolean`                       | If true the method won't block. Defaults to false. If the buffer is modified before the formatter completes, the formatting will be discarded.       |
 |          | dry_run                                              | `nil\|boolean`                       | If true don't apply formatting changes to the buffer                                                                                                 |
+|          | undojoin                                             | `nil\|boolean`                       | Use undojoin to merge formatting changes with previous edit (default false)                                                                          |
 |          | formatters                                           | `nil\|string[]`                      | List of formatters to run. Defaults to all formatters for the buffer filetype.                                                                       |
 |          | lsp_format                                           | `nil\|conform.LspFormatOpts`         | Configure if and when LSP should be used for formatting. Defaults to "never".                                                                        |
 |          |                                                      | > `"never"`                          | never use the LSP for formatting (default)                                                                                                           |
diff --git a/doc/conform.txt b/doc/conform.txt
index e27d9b7..3f6568f 100644
--- a/doc/conform.txt
+++ b/doc/conform.txt
@@ -150,6 +150,8 @@ format({opts}, {callback}): boolean                               *conform.forma
                        completes, the formatting will be discarded.
           {dry_run}    `nil|boolean` If true don't apply formatting changes to
                        the buffer
+          {undojoin}   `nil|boolean` Use undojoin to merge formatting changes
+                       with previous edit (default false)
           {formatters} `nil|string[]` List of formatters to run. Defaults to all
                        formatters for the buffer filetype.
           {lsp_format} `nil|conform.LspFormatOpts` Configure if and when LSP
diff --git a/lua/conform/init.lua b/lua/conform/init.lua
index aa97f1c..f0b2dfd 100644
--- a/lua/conform/init.lua
+++ b/lua/conform/init.lua
@@ -340,6 +340,7 @@ end
 ---@field bufnr nil|integer Format this buffer (default 0)
 ---@field async nil|boolean If true the method won't block. Defaults to false. If the buffer is modified before the formatter completes, the formatting will be discarded.
 ---@field dry_run nil|boolean If true don't apply formatting changes to the buffer
+---@field undojoin nil|boolean Use undojoin to merge formatting changes with previous edit (default false)
 ---@field formatters nil|string[] List of formatters to run. Defaults to all formatters for the buffer filetype.
 ---@field lsp_format? conform.LspFormatOpts Configure if and when LSP should be used for formatting. Defaults to "never".
 ---@field quiet nil|boolean Don't show any notifications for warnings or failures. Defaults to false.
@@ -353,7 +354,7 @@ end
 ---@param callback? fun(err: nil|string, did_edit: nil|boolean) Called once formatting has completed
 ---@return boolean True if any formatters were attempted
 M.format = function(opts, callback)
-  ---@type {timeout_ms: integer, bufnr: integer, async: boolean, dry_run: boolean, lsp_format: "never"|"first"|"last"|"prefer"|"fallback", quiet: boolean, formatters?: string[], range?: conform.Range}
+  ---@type {timeout_ms: integer, bufnr: integer, async: boolean, dry_run: boolean, lsp_format: "never"|"first"|"last"|"prefer"|"fallback", quiet: boolean, formatters?: string[], range?: conform.Range, undojoin: boolean}
   opts = vim.tbl_extend("keep", opts or {}, {
     timeout_ms = 1000,
     bufnr = 0,
@@ -361,6 +362,7 @@ M.format = function(opts, callback)
     dry_run = false,
     lsp_format = "never",
     quiet = false,
+    undojoin = false,
   })
 
   -- For backwards compatibility
@@ -430,7 +432,8 @@ M.format = function(opts, callback)
       return f.name
     end, formatters)
     log.debug("Running formatters on %s: %s", vim.api.nvim_buf_get_name(opts.bufnr), resolved_names)
-    local run_opts = { exclusive = true, dry_run = opts.dry_run }
+    ---@type conform.RunOpts
+    local run_opts = { exclusive = true, dry_run = opts.dry_run, undojoin = opts.undojoin }
     if opts.async then
       runner.format_async(opts.bufnr, formatters, opts.range, run_opts, cb)
     else
@@ -517,7 +520,8 @@ M.format_lines = function(formatter_names, lines, opts, callback)
     callback(err, new_lines)
   end
 
-  local run_opts = { exclusive = false, dry_run = false }
+  ---@type conform.RunOpts
+  local run_opts = { exclusive = false, dry_run = false, undojoin = false }
   if opts.async then
     runner.format_lines_async(opts.bufnr, formatters, nil, lines, run_opts, handle_err)
   else
diff --git a/lua/conform/lsp_format.lua b/lua/conform/lsp_format.lua
index a4f9556..fe26656 100644
--- a/lua/conform/lsp_format.lua
+++ b/lua/conform/lsp_format.lua
@@ -4,7 +4,7 @@ local util = require("vim.lsp.util")
 
 local M = {}
 
-local function apply_text_edits(text_edits, bufnr, offset_encoding, dry_run)
+local function apply_text_edits(text_edits, bufnr, offset_encoding, dry_run, undojoin)
   if
     #text_edits == 1
     and text_edits[1].range.start.line == 0
@@ -25,11 +25,15 @@ local function apply_text_edits(text_edits, bufnr, offset_encoding, dry_run)
       new_lines,
       nil,
       false,
-      dry_run
+      dry_run,
+      undojoin
     )
   elseif dry_run then
     return #text_edits > 0
   else
+    if undojoin then
+      vim.cmd.undojoin()
+    end
     vim.lsp.util.apply_text_edits(text_edits, bufnr, offset_encoding)
     return #text_edits > 0
   end
@@ -124,8 +128,13 @@ function M.format(options, callback)
             )
           )
         else
-          local this_did_edit =
-            apply_text_edits(result, ctx.bufnr, client.offset_encoding, options.dry_run)
+          local this_did_edit = apply_text_edits(
+            result,
+            ctx.bufnr,
+            client.offset_encoding,
+            options.dry_run,
+            options.undojoin
+          )
           changedtick = vim.b[bufnr].changedtick
 
           if options.dry_run and this_did_edit then
@@ -145,8 +154,13 @@ function M.format(options, callback)
       local params = set_range(client, util.make_formatting_params(options.formatting_options))
       local result, err = client.request_sync(method, params, timeout_ms, bufnr)
       if result and result.result then
-        local this_did_edit =
-          apply_text_edits(result.result, bufnr, client.offset_encoding, options.dry_run)
+        local this_did_edit = apply_text_edits(
+          result.result,
+          bufnr,
+          client.offset_encoding,
+          options.dry_run,
+          options.undojoin
+        )
         did_edit = did_edit or this_did_edit
 
         if options.dry_run and did_edit then
diff --git a/lua/conform/runner.lua b/lua/conform/runner.lua
index c0a19cd..333c16f 100644
--- a/lua/conform/runner.lua
+++ b/lua/conform/runner.lua
@@ -9,6 +9,7 @@ local M = {}
 ---@class (exact) conform.RunOpts
 ---@field exclusive boolean If true, ensure only a single formatter is running per buffer
 ---@field dry_run boolean If true, do not apply changes and stop after the first formatter attempts to do so
+---@field undojoin boolean Use undojoin to merge formatting changes with previous edit
 
 ---@param formatter_name string
 ---@param ctx conform.Context
@@ -168,8 +169,18 @@ end
 ---@param new_lines string[]
 ---@param range? conform.Range
 ---@param only_apply_range boolean
+---@param dry_run boolean
+---@param undojoin boolean
 ---@return boolean any_changes
-M.apply_format = function(bufnr, original_lines, new_lines, range, only_apply_range, dry_run)
+M.apply_format = function(
+  bufnr,
+  original_lines,
+  new_lines,
+  range,
+  only_apply_range,
+  dry_run,
+  undojoin
+)
   if bufnr == 0 then
     bufnr = vim.api.nvim_get_current_buf()
   end
@@ -251,6 +262,9 @@ M.apply_format = function(bufnr, original_lines, new_lines, range, only_apply_ra
 
   if not dry_run then
     log.trace("Applying text edits: %s", text_edits)
+    if undojoin then
+      vim.cmd.undojoin()
+    end
     vim.lsp.util.apply_text_edits(text_edits, bufnr, "utf-8")
     log.trace("Done formatting %s", bufname)
   end
@@ -535,7 +549,8 @@ M.format_async = function(bufnr, formatters, range, opts, callback)
           output_lines,
           range,
           not all_support_range_formatting,
-          opts.dry_run
+          opts.dry_run,
+          opts.undojoin
         )
       end
       callback(err, did_edit)
@@ -609,7 +624,8 @@ M.format_sync = function(bufnr, formatters, timeout_ms, range, opts)
     final_result,
     range,
     not all_support_range_formatting,
-    opts.dry_run
+    opts.dry_run,
+    opts.undojoin
   )
   return err, did_edit
 end
diff --git a/tests/fuzzer_spec.lua b/tests/fuzzer_spec.lua
index c47f1e5..b110f75 100644
--- a/tests/fuzzer_spec.lua
+++ b/tests/fuzzer_spec.lua
@@ -25,7 +25,7 @@ describe("fuzzer", function()
     vim.api.nvim_set_current_buf(bufnr)
     vim.api.nvim_buf_set_lines(bufnr, 0, -1, true, buf_content)
     vim.bo[bufnr].modified = false
-    runner.apply_format(0, buf_content, expected, nil, false)
+    runner.apply_format(0, buf_content, expected, nil, false, false, false)
     assert.are.same(expected, vim.api.nvim_buf_get_lines(0, 0, -1, false))
   end