feat: Repeat delimiters for add.

TODO: Repeat deletions, unit test stuff.

Author: kylechui

lua/nvim-surround/annotations.lua

@@ -19,6 +19,11 @@
 ---@field left selection|nil
 ---@field right selection|nil
 
+-- TODO: Figure out a better name for this, since it's stupid for get_input to NOT return user_input
+---@class user_input
+---@field char string|nil
+---@field count integer
+
 --[====================================================================================================================[
                                                     Internal Options
 --]====================================================================================================================]

lua/nvim-surround/config.lua

@@ -350,12 +350,13 @@ M.get_alias = function(char)
 end
 
 -- Gets a delimiter pair for a user-inputted character.
----@param char string|nil The user-given character.
+---@param user_input user_input The user-given character, and number of times to repeat the delimiters.
 ---@param line_mode boolean Whether or not the delimiters should be put on new lines.
 ---@return delimiter_pair|nil @A pair of delimiters for the given input, or nil if not applicable.
 ---@nodiscard
-M.get_delimiters = function(char, line_mode)
-    char = M.get_alias(char)
+M.get_delimiters = function(user_input, line_mode)
+    local utils = require("nvim-surround.utils")
+    local char = M.get_alias(user_input.char)
     -- Get the delimiters, using invalid_key_behavior if the add function is undefined for the character
     local delimiters = M.get_add(char)(char)
     if delimiters == nil then
@@ -378,7 +379,7 @@ M.get_delimiters = function(char, line_mode)
         table.insert(lhs, "")
     end
 
-    return { lhs, rhs }
+    return utils.repeat_delimiters({ lhs, rhs }, user_input.count)
 end
 
 -- Returns the add key for the surround associated with a given character, if one exists.
@@ -547,8 +548,7 @@ M.translate_opts = function(user_opts)
     local input = require("nvim-surround.input")
     local opts = {}
     for key, value in pairs(user_opts) do
-        if key == "surrounds" then
-        elseif key == "indent_lines" then
+        if key == "indent_lines" then
             opts[key] = value or function() end
         else
             opts[key] = value

lua/nvim-surround/init.lua

@@ -328,7 +328,9 @@ M.delete_callback = function()
     -- Save the current position of the cursor
     local curpos = buffer.get_curpos()
     -- Get a character input if not cached
-    cache.delete.char = cache.delete.char or input.get_char()
+    -- TODO: repeat deletions!
+    local user_input = input.get_char()
+    cache.delete.char = cache.delete.char or user_input.char
     if not cache.delete.char then
         return
     end
@@ -348,7 +350,8 @@ M.change_callback = function()
     -- Save the current position of the cursor
     local curpos = buffer.get_curpos()
     if not cache.change.del_char or not cache.change.add_delimiters then
-        local del_char = config.get_alias(input.get_char())
+        local user_input = input.get_char()
+        local del_char = config.get_alias(user_input.char)
         local change = config.get_change(del_char)
         local selections = utils.get_nearest_selections(del_char, "change")
         if not (del_char and change and selections) then

lua/nvim-surround/input.lua

@@ -14,15 +14,31 @@ M.replace_termcodes = function(char)
 end
 
 -- Gets a character input from the user.
----@return string|nil @The input character, or nil if an escape character is pressed.
+---@return user_input @The input character, or nil if an escape character is pressed, with optional count.
 ---@nodiscard
 M.get_char = function()
-    local ok, char = pcall(vim.fn.getcharstr)
-    -- Return nil if input is cancelled (e.g. <C-c> or <Esc>)
-    if not ok or char == "\27" then
-        return nil
-    end
-    return M.replace_termcodes(char)
+    local output = { count = nil, char = nil }
+
+    repeat
+        local ok, char = pcall(vim.fn.getcharstr)
+        -- Return nil if input is cancelled (e.g. <C-c> or <Esc>)
+        if not ok or char == "\27" then
+            return { count = 0, char = nil }
+        end
+
+        local digit = tonumber(char)
+        if digit then
+            output.count = output.count == nil and 0 or output.count
+            output.count = output.count * 10 + digit
+        else
+            output.char = M.replace_termcodes(char)
+        end
+    until output.char ~= nil
+
+    return {
+        count = output.count == nil and 1 or output.count,
+        char = output.char,
+    }
 end
 
 -- Gets a string input from the user.

lua/nvim-surround/utils.lua

@@ -7,6 +7,21 @@ local M = {}
 -- Do nothing.
 M.NOOP = function() end
 
+-- Repeats a delimiter pair n times.
+---@param delimiters delimiter_pair The delimiters to be repeated.
+---@param n integer The number of times to repeat the delimiters.
+---@nodiscard
+M.repeat_delimiters = function(delimiters, n)
+    local acc = { { "" }, { "" } }
+    for _ = 1, n do
+        acc[1][#acc[1]] = acc[1][#acc[1]] .. delimiters[1][1]
+        vim.list_extend(acc[1], delimiters[1], 2)
+        acc[2][#acc[2]] = acc[2][#acc[2]] .. delimiters[2][1]
+        vim.list_extend(acc[2], delimiters[2], 2)
+    end
+    return acc
+end
+
 -- Gets the nearest two selections for the left and right surrounding pair.
 ---@param char string|nil A character representing what kind of surrounding pair is to be selected.
 ---@param action "delete"|"change" A string representing what action is being performed.