feat(input): support for start_mode (#175)

* Add support for start_mode

* Document options for start_mode

* Add back restore mode

* Move shims to config.lua

* Separate set_mode and restore_mode

* Fix select mode

* Fix luacheck and stylua

* Fix typecheck

* Add tests for start_mode

* Fix stylua

* Implement review comment about restore_mode

---------

Co-authored-by: Marcel Robitaille <[email protected]>

README.md

@@ -141,8 +141,8 @@ require("dressing").setup({
     -- Can be 'left', 'right', or 'center'
     title_pos = "left",
 
-    -- When true, input will start in insert mode.
-    start_in_insert = true,
+    -- The initial mode when the window opens (insert|normal|visual|select).
+    start_mode = "insert",
 
     -- These are passed to nvim_open_win
     border = "rounded",

doc/dressing.txt

@@ -19,8 +19,8 @@ Configure dressing.nvim by calling the setup() function.
         -- Can be 'left', 'right', or 'center'
         title_pos = "left",
 
-        -- When true, input will start in insert mode.
-        start_in_insert = true,
+        -- The initial mode when the window opens (insert|normal|visual|select).
+        start_mode = "insert",
 
         -- These are passed to nvim_open_win
         border = "rounded",

lua/dressing/config.lua

@@ -12,8 +12,8 @@ local default_config = {
     -- Can be 'left', 'right', or 'center'
     title_pos = "left",
 
-    -- When true, input will start in insert mode.
-    start_in_insert = true,
+    -- The initial mode when the window opens (insert|normal|visual|select).
+    start_mode = "insert",
 
     -- These are passed to nvim_open_win
     border = "rounded",
@@ -163,11 +163,24 @@ local default_config = {
 
 local M = vim.deepcopy(default_config)
 
+-- Apply shims for backwards compatibility
+---@param key string
+---@param opts table
+---@return table
+M.apply_shim = function(key, opts)
+  -- Support start_in_insert for backwards compatibility.
+  if key == "input" and opts.start_in_insert ~= nil then
+    opts.start_mode = opts.start_in_insert and "insert" or "normal"
+  end
+
+  return opts
+end
+
 M.update = function(opts)
   local newconf = vim.tbl_deep_extend("force", default_config, opts or {})
 
   for k, v in pairs(newconf) do
-    M[k] = v
+    M[k] = M.apply_shim(k, v)
   end
 end
 
@@ -177,7 +190,10 @@ M.get_mod_config = function(key, ...)
   if not M[key].get_config then
     return M[key]
   end
+
   local conf = M[key].get_config(...)
+  conf = M.apply_shim(key, conf)
+
   if conf then
     return vim.tbl_deep_extend("force", M[key], conf)
   else

lua/dressing/input.lua

@@ -4,13 +4,26 @@ local patch = require("dressing.patch")
 local util = require("dressing.util")
 local M = {}
 
+---@alias dressing.Mode "insert" | "visual" | "normal" | "select"
+
 ---@class (exact) dressing.InputContext
 ---@field opts? dressing.InputOptions
 ---@field on_confirm? fun(text?: string)
 ---@field winid? integer
 ---@field history_idx? integer
 ---@field history_tip? string
+---@field mode_to_restore? dressing.Mode The mode to restore when the input window closes.
+
+---@class (exact) dressing.InputConfig
 ---@field start_in_insert? boolean
+---@field start_mode? dressing.Mode
+---@field enabled? boolean
+---@field default_prompt? string
+---@field win_options? table
+---@field buf_options? table
+---@field mappings? table
+---@field trim_prompt? boolean
+---@field prompt_align? string
 
 ---@class (exact) dressing.InputOptions
 ---@field prompt? string
@@ -26,7 +39,7 @@ local context = {
   winid = nil,
   history_idx = nil,
   history_tip = nil,
-  start_in_insert = nil,
+  mode_to_restore = nil,
 }
 
 local keymaps = {
@@ -91,6 +104,27 @@ M.history_next = function()
   end
 end
 
+---@param mode dressing.Mode?
+M.set_mode = function(mode)
+  if mode == "normal" then
+    vim.cmd("stopinsert")
+  elseif mode == "insert" then
+    vim.cmd("startinsert!")
+  elseif mode == "visual" then
+    vim.api.nvim_command("normal! vg_")
+  elseif mode == "select" then
+    M.set_mode("visual")
+    vim.api.nvim_feedkeys(vim.api.nvim_replace_termcodes("<C-g>", true, false, true), "n", true)
+  end
+end
+
+---@param mode dressing.Mode?
+M.restore_mode = function(mode)
+  if mode ~= "insert" then
+    vim.cmd("stopinsert")
+  end
+end
+
 local function close_completion_window()
   if vim.fn.pumvisible() == 1 then
     local escape_key = vim.api.nvim_replace_termcodes("<C-e>", true, false, true)
@@ -105,9 +139,8 @@ local function confirm(text)
   close_completion_window()
   local ctx = context
   context = {}
-  if not ctx.start_in_insert then
-    vim.cmd("stopinsert")
-  end
+  M.restore_mode(ctx.mode_to_restore)
+
   -- We have to wait briefly for the popup window to close (if present),
   -- otherwise vim gets into a very weird and bad state. I was seeing text get
   -- deleted from the buffer after the input window closes.
@@ -266,7 +299,7 @@ end
 ---@param prompt_lines string[]
 ---@param default? string
 ---@return integer
----@return boolean
+---@return dressing.Mode?
 local function create_or_update_win(config, prompt_lines, default)
   local parent_win = 0
   local winopt
@@ -326,18 +359,19 @@ local function create_or_update_win(config, prompt_lines, default)
 
   winopt = config.override(winopt) or winopt
 
-  local winid, start_in_insert
+  local winid, mode_to_restore
   -- If the floating win was already open
   if win_conf then
     -- Make sure the previous on_confirm callback is called with nil
     vim.schedule(context.on_confirm)
     vim.api.nvim_win_set_config(context.winid, winopt)
     winid = context.winid
-    start_in_insert = context.start_in_insert
+    mode_to_restore = context.mode_to_restore
   else
-    start_in_insert = string.sub(vim.api.nvim_get_mode().mode, 1, 1) == "i"
     local bufnr = vim.api.nvim_create_buf(false, true)
     winid = vim.api.nvim_open_win(bufnr, true, winopt)
+    local mode_chr = string.sub(vim.api.nvim_get_mode().mode, 1, 1)
+    mode_to_restore = ({ i = "insert", n = "normal", v = "visual", s = "select" })[mode_chr]
   end
 
   -- If the prompt is multiple lines, create another window for it
@@ -387,8 +421,7 @@ local function create_or_update_win(config, prompt_lines, default)
   end
 
   ---@cast winid integer
-  ---@cast start_in_insert boolean
-  return winid, start_in_insert
+  return winid, mode_to_restore
 end
 
 ---@param opts string|dressing.InputOptions
@@ -401,22 +434,24 @@ local show_input = util.make_queued_async_fn(2, function(opts, on_confirm)
   if type(opts) ~= "table" then
     opts = { prompt = tostring(opts) }
   end
-  local config = global_config.get_mod_config("input", opts)
+  local config = global_config.get_mod_config("input", opts) --[[@as dressing.InputConfig]]
   if not config.enabled then
     return patch.original_mods.input(opts, on_confirm)
   end
 
-  local prompt = opts.prompt or config.default_prompt
+  local start_mode = config.start_mode
+
+  local prompt = opts.prompt or config.default_prompt --[[@as string]]
   local prompt_lines = vim.split(prompt, "\n", { plain = true, trimempty = true })
 
   -- Create or update the window
-  local winid, start_in_insert = create_or_update_win(config, prompt_lines, opts.default)
+  local winid, mode_to_restore = create_or_update_win(config, prompt_lines, opts.default)
   context = {
     winid = winid,
     on_confirm = on_confirm,
     opts = opts,
     history_idx = nil,
-    start_in_insert = start_in_insert,
+    mode_to_restore = mode_to_restore,
   }
   for option, value in pairs(config.win_options) do
     vim.api.nvim_set_option_value(option, value, { scope = "local", win = winid })
@@ -484,9 +519,9 @@ local show_input = util.make_queued_async_fn(2, function(opts, on_confirm)
     callback = M.close,
   })
 
-  if config.start_in_insert then
-    vim.cmd("startinsert!")
-  end
+  ---@cast start_mode dressing.Mode
+  M.set_mode(start_mode)
+
   close_completion_window()
   apply_highlight()
 end)

tests/input_spec.lua

@@ -1,7 +1,11 @@
 require("plenary.async").tests.add_to_env()
+local a = require("plenary.async")
 local dressing = require("dressing")
+local input = require("dressing.input")
 local util = require("tests.util")
 local channel = a.control.channel
+local assert = require("luassert")
+local stub = require("luassert.stub")
 
 local function run_input(keys, opts)
   opts = opts or {}
@@ -81,28 +85,6 @@ a.describe("input modal", function()
     assert(ret == "", string.format("Got '%s' expected nil", ret))
   end)
 
-  a.it("starts in normal mode when start_in_insert = false", function()
-    local orig_cmd = vim.cmd
-    local startinsert_called = false
-    vim.cmd = function(cmd)
-      if cmd == "startinsert!" then
-        startinsert_called = true
-      end
-      orig_cmd(cmd)
-    end
-
-    require("dressing.config").input.start_in_insert = false
-    run_input({
-      "my text",
-      "<CR>",
-    }, {
-      after_fn = function()
-        vim.cmd = orig_cmd
-      end,
-    })
-    assert(not startinsert_called, "Got 'true' expected 'false'")
-  end)
-
   a.it("queues successive calls to vim.ui.input", function()
     local tx1, rx1 = channel.oneshot()
     local tx2, rx2 = channel.oneshot()
@@ -139,6 +121,84 @@ a.describe("input modal", function()
     assert(vim.fn.pumvisible() == 0, "Popup menu should not be visible after leaving modal")
   end)
 
+  local function test_start_mode(expected_start_mode, expected_restore_mode)
+    -- Since going into insert mode does not work well in headless, use mocks.
+    local nvim_get_mode = stub(vim.api, "nvim_get_mode", { mode = expected_restore_mode })
+    local set_mode = stub(input, "set_mode")
+    local restore_mode = stub(input, "restore_mode")
+
+    local tx, rx = channel.oneshot()
+    vim.ui.input({}, tx)
+
+    assert.stub(set_mode).was_called(1)
+    assert.spy(set_mode).was_called_with(expected_start_mode)
+    set_mode:clear()
+    assert.stub(restore_mode).was_not_called()
+
+    util.feedkeys({ "<CR>" })
+
+    assert.spy(set_mode).was_not_called()
+    assert.stub(restore_mode).was_called(1)
+    assert.stub(restore_mode).was_called_with(expected_restore_mode)
+
+    rx()
+
+    set_mode:revert()
+    restore_mode:revert()
+    nvim_get_mode:revert()
+  end
+
+  -- Visual causes problems. The other's should be enough. The logic is the same.
+  for _, start_mode in ipairs({ "normal", "insert", "select" }) do
+    -- Only normal and insert are supported.
+    for _, mode_to_restore in ipairs({ "normal", "insert" }) do
+      a.it(
+        "sets the mode correctly to start_mode="
+          .. start_mode
+          .. " restore_mode="
+          .. mode_to_restore,
+        function()
+          require("dressing.config").input.start_mode = start_mode
+          test_start_mode(start_mode, mode_to_restore)
+        end
+      )
+    end
+  end
+
+  a.it("is backwards compatible with start_in_insert = false", function()
+    require("dressing.config").update({ input = { start_in_insert = false } })
+    test_start_mode("normal", "normal")
+  end)
+
+  a.it("is backwards compatible with start_in_insert = true", function()
+    require("dressing.config").update({ input = { start_in_insert = true } })
+    test_start_mode("insert", "normal")
+  end)
+
+  a.it("get_config takes precedence (normal)", function()
+    require("dressing.config").update({
+      input = {
+        start_mode = "insert",
+        get_config = function()
+          return { start_in_insert = false }
+        end,
+      },
+    })
+    test_start_mode("normal", "normal")
+  end)
+
+  a.it("get_config takes precedence (insert)", function()
+    require("dressing.config").update({
+      input = {
+        start_mode = "normal",
+        get_config = function()
+          return { start_in_insert = true }
+        end,
+      },
+    })
+    test_start_mode("insert", "normal")
+  end)
+
   a.it("can cancel out when popup menu is open", function()
     vim.opt.completeopt = { "menu", "menuone", "noselect" }
     local ret = run_input({