Merge pull request #20387 from famiu/feat/nvim_cmd/buffer-args

feat(nvim_cmd)!: allow using first argument as count

4 files changed, 117 insertions, 30 deletions

diff --git a/runtime/doc/api.txt b/runtime/doc/api.txt
index 59d1d36be9..9515dea9a2 100644
--- a/runtime/doc/api.txt
+++ b/runtime/doc/api.txt
@@ -1697,7 +1697,13 @@ nvim_cmd({*cmd}, {*opts})                                         *nvim_cmd()*
     of a String. This allows for easier construction and manipulation of an Ex
     command. This also allows for things such as having spaces inside a
     command argument, expanding filenames in a command that otherwise doesn't
-    expand filenames, etc.
+    expand filenames, etc. Command arguments may also be Number, Boolean or
+    String.
+
+    The first argument may also be used instead of count for commands that
+    support it in order to make their usage simpler with |vim.cmd()|. For
+    example, instead of `vim.cmd.bdelete{ count = 2 }`, you may do
+    `vim.cmd.bdelete(2)`.
 
     On execution error: fails with VimL error, updates v:errmsg.
 
diff --git a/src/nvim/api/command.c b/src/nvim/api/command.c
index ab166c6b38..d410466c44 100644
--- a/src/nvim/api/command.c
+++ b/src/nvim/api/command.c
@@ -283,7 +283,11 @@ end:
 /// Unlike |nvim_command()| this command takes a structured Dictionary instead of a String. This
 /// allows for easier construction and manipulation of an Ex command. This also allows for things
 /// such as having spaces inside a command argument, expanding filenames in a command that otherwise
-/// doesn't expand filenames, etc.
+/// doesn't expand filenames, etc. Command arguments may also be Number, Boolean or String.
+///
+/// The first argument may also be used instead of count for commands that support it in order to
+/// make their usage simpler with |vim.cmd()|. For example, instead of
+/// `vim.cmd.bdelete{ count = 2 }`, you may do `vim.cmd.bdelete(2)`.
 ///
 /// On execution error: fails with VimL error, updates v:errmsg.
 ///
@@ -308,8 +312,7 @@ String nvim_cmd(uint64_t channel_id, Dict(cmd) *cmd, Dict(cmd_opts) *opts, Error
 
   char *cmdline = NULL;
   char *cmdname = NULL;
-  ArrayOf(String) args;
-  size_t argc = 0;
+  ArrayOf(String) args = ARRAY_DICT_INIT;
 
   String retv = (String)STRING_INIT;
 
@@ -381,48 +384,70 @@ String nvim_cmd(uint64_t channel_id, Dict(cmd) *cmd, Dict(cmd_opts) *opts, Error
     if (cmd->args.type != kObjectTypeArray) {
       VALIDATION_ERROR("'args' must be an Array");
     }
-    // Check if every argument is valid
+
+    // Process all arguments. Convert non-String arguments to String and check if String arguments
+    // have non-whitespace characters.
     for (size_t i = 0; i < cmd->args.data.array.size; i++) {
       Object elem = cmd->args.data.array.items[i];
-      if (elem.type != kObjectTypeString) {
-        VALIDATION_ERROR("Command argument must be a String");
-      } else if (string_iswhite(elem.data.string)) {
-        VALIDATION_ERROR("Command argument must have non-whitespace characters");
+      char *data_str;
+
+      switch (elem.type) {
+      case kObjectTypeBoolean:
+        data_str = xcalloc(2, sizeof(char));
+        data_str[0] = elem.data.boolean ? '1' : '0';
+        data_str[1] = '\0';
+        break;
+      case kObjectTypeBuffer:
+      case kObjectTypeWindow:
+      case kObjectTypeTabpage:
+      case kObjectTypeInteger:
+        data_str = xcalloc(NUMBUFLEN, sizeof(char));
+        snprintf(data_str, NUMBUFLEN, "%" PRId64, elem.data.integer);
+        break;
+      case kObjectTypeString:
+        if (string_iswhite(elem.data.string)) {
+          VALIDATION_ERROR("String command argument must have at least one non-whitespace "
+                           "character");
+        }
+        data_str = xstrndup(elem.data.string.data, elem.data.string.size);
+        break;
+      default:
+        VALIDATION_ERROR("Invalid type for command argument");
+        break;
       }
+
+      ADD(args, STRING_OBJ(cstr_as_string(data_str)));
     }
 
-    argc = cmd->args.data.array.size;
     bool argc_valid;
 
     // Check if correct number of arguments is used.
     switch (ea.argt & (EX_EXTRA | EX_NOSPC | EX_NEEDARG)) {
     case EX_EXTRA | EX_NOSPC | EX_NEEDARG:
-      argc_valid = argc == 1;
+      argc_valid = args.size == 1;
       break;
     case EX_EXTRA | EX_NOSPC:
-      argc_valid = argc <= 1;
+      argc_valid = args.size <= 1;
       break;
     case EX_EXTRA | EX_NEEDARG:
-      argc_valid = argc >= 1;
+      argc_valid = args.size >= 1;
       break;
     case EX_EXTRA:
       argc_valid = true;
       break;
     default:
-      argc_valid = argc == 0;
+      argc_valid = args.size == 0;
       break;
     }
 
     if (!argc_valid) {
       VALIDATION_ERROR("Incorrect number of arguments supplied");
     }
-
-    args = cmd->args.data.array;
   }
 
   // Simply pass the first argument (if it exists) as the arg pointer to `set_cmd_addr_type()`
   // since it only ever checks the first argument.
-  set_cmd_addr_type(&ea, argc > 0 ? args.items[0].data.string.data : NULL);
+  set_cmd_addr_type(&ea, args.size > 0 ? args.items[0].data.string.data : NULL);
 
   if (HAS_KEY(cmd->range)) {
     if (!(ea.argt & EX_RANGE)) {
@@ -626,7 +651,7 @@ String nvim_cmd(uint64_t channel_id, Dict(cmd) *cmd, Dict(cmd_opts) *opts, Error
 
   // Finally, build the command line string that will be stored inside ea.cmdlinep.
   // This also sets the values of ea.cmd, ea.arg, ea.args and ea.arglens.
-  build_cmdline_str(&cmdline, &ea, &cmdinfo, args, argc);
+  build_cmdline_str(&cmdline, &ea, &cmdinfo, args);
   ea.cmdlinep = &cmdline;
 
   garray_T capture_local;
@@ -682,6 +707,7 @@ clear_ga:
     ga_clear(&capture_local);
   }
 end:
+  api_free_array(args);
   xfree(cmdline);
   xfree(cmdname);
   xfree(ea.args);
@@ -711,8 +737,9 @@ static bool string_iswhite(String str)
 
 /// Build cmdline string for command, used by `nvim_cmd()`.
 static void build_cmdline_str(char **cmdlinep, exarg_T *eap, CmdParseInfo *cmdinfo,
-                              ArrayOf(String) args, size_t argc)
+                              ArrayOf(String) args)
 {
+  size_t argc = args.size;
   StringBuilder cmdline = KV_INITIAL_VALUE;
   kv_resize(cmdline, 32);  // Make it big enough to handle most typical commands
 
@@ -798,7 +825,7 @@ static void build_cmdline_str(char **cmdlinep, exarg_T *eap, CmdParseInfo *cmdin
   }
 
   eap->argc = argc;
-  eap->arglens = xcalloc(argc, sizeof(size_t));
+  eap->arglens = eap->argc > 0 ? xcalloc(argc, sizeof(size_t)) : NULL;
   size_t argstart_idx = cmdline.size;
   for (size_t i = 0; i < argc; i++) {
     String s = args.items[i].data.string;
@@ -813,7 +840,7 @@ static void build_cmdline_str(char **cmdlinep, exarg_T *eap, CmdParseInfo *cmdin
   // Now that all the arguments are appended, use the command index and argument indices to set the
   // values of eap->cmd, eap->arg and eap->args.
   eap->cmd = cmdline.items + cmdname_idx;
-  eap->args = xcalloc(argc, sizeof(char *));
+  eap->args = eap->argc > 0 ? xcalloc(argc, sizeof(char *)) : NULL;
   size_t offset = argstart_idx;
   for (size_t i = 0; i < argc; i++) {
     offset++;  // Account for space
diff --git a/src/nvim/ex_docmd.c b/src/nvim/ex_docmd.c
index 386f708f79..987fb439d7 100644
--- a/src/nvim/ex_docmd.c
+++ b/src/nvim/ex_docmd.c
@@ -1339,6 +1339,20 @@ static int parse_count(exarg_T *eap, char **errormsg, bool validate)
           || ascii_iswhite(*p))) {
     long n = getdigits_long(&eap->arg, false, -1);
     eap->arg = skipwhite(eap->arg);
+
+    if (eap->args != NULL) {
+      assert(eap->argc > 0 && eap->arg >= eap->args[0]);
+      // If eap->arg is still pointing to the first argument, just make eap->args[0] point to the
+      // same location. This is needed for usecases like vim.cmd.sleep('10m'). If eap->arg is
+      // pointing outside the first argument, shift arguments by 1.
+      if (eap->arg < eap->args[0] + eap->arglens[0]) {
+        eap->arglens[0] -= (size_t)(eap->arg - eap->args[0]);
+        eap->args[0] = eap->arg;
+      } else {
+        shift_cmd_args(eap);
+      }
+    }
+
     if (n <= 0 && (eap->argt & EX_ZEROR) == 0) {
       if (errormsg != NULL) {
         *errormsg = _(e_zerocount);
@@ -1512,6 +1526,30 @@ end:
   return retval;
 }
 
+// Shift Ex-command arguments to the right.
+static void shift_cmd_args(exarg_T *eap)
+{
+  assert(eap->args != NULL && eap->argc > 0);
+
+  char **oldargs = eap->args;
+  size_t *oldarglens = eap->arglens;
+
+  eap->argc--;
+  eap->args = eap->argc > 0 ? xcalloc(eap->argc, sizeof(char *)) : NULL;
+  eap->arglens = eap->argc > 0 ? xcalloc(eap->argc, sizeof(size_t)) : NULL;
+
+  for (size_t i = 0; i < eap->argc; i++) {
+    eap->args[i] = oldargs[i + 1];
+    eap->arglens[i] = oldarglens[i + 1];
+  }
+
+  // If there are no arguments, make eap->arg point to the end of string.
+  eap->arg = (eap->argc > 0 ? eap->args[0] : (oldargs[0] + oldarglens[0]));
+
+  xfree(oldargs);
+  xfree(oldarglens);
+}
+
 static int execute_cmd0(int *retv, exarg_T *eap, char **errormsg, bool preview)
 {
   // If filename expansion is enabled, expand filenames
@@ -1551,15 +1589,7 @@ static int execute_cmd0(int *retv, exarg_T *eap, char **errormsg, bool preview)
                                    eap->args[0] + eap->arglens[0],
                                    (eap->argt & EX_BUFUNL) != 0, false, false);
       eap->addr_count = 1;
-      // Shift each argument by 1
-      for (size_t i = 0; i < eap->argc - 1; i++) {
-        eap->args[i] = eap->args[i + 1];
-      }
-      // Make the last argument point to the NUL terminator at the end of string
-      eap->args[eap->argc - 1] = eap->args[eap->argc - 1] + eap->arglens[eap->argc - 1];
-      eap->argc -= 1;
-
-      eap->arg = eap->args[0];
+      shift_cmd_args(eap);
     }
     if (eap->line2 < 0) {  // failed
       return FAIL;
@@ -1658,6 +1688,11 @@ int execute_cmd(exarg_T *eap, CmdParseInfo *cmdinfo, bool preview)
     (void)hasFolding(eap->line2, NULL, &eap->line2);
   }
 
+  // Use first argument as count when possible
+  if (parse_count(eap, &errormsg, true) == FAIL) {
+    goto end;
+  }
+
   // Execute the command
   execute_cmd0(&retv, eap, &errormsg, preview);
 
diff --git a/test/functional/api/vim_spec.lua b/test/functional/api/vim_spec.lua
index c2f3a5ec5e..ca1c5070a6 100644
--- a/test/functional/api/vim_spec.lua
+++ b/test/functional/api/vim_spec.lua
@@ -1,6 +1,7 @@
 local helpers = require('test.functional.helpers')(after_each)
 local Screen = require('test.functional.ui.screen')
 local lfs = require('lfs')
+local luv = require('luv')
 
 local fmt = string.format
 local assert_alive = helpers.assert_alive
@@ -3962,5 +3963,23 @@ describe('API', function()
         15                                      |
       ]]}
     end)
+    it('works with non-String args', function()
+      eq('2', meths.cmd({cmd = 'echo', args = {2}}, {output = true}))
+      eq('1', meths.cmd({cmd = 'echo', args = {true}}, {output = true}))
+    end)
+    describe('first argument as count', function()
+      before_each(clear)
+
+      it('works', function()
+        command('vsplit | enew')
+        meths.cmd({cmd = 'bdelete', args = {meths.get_current_buf()}}, {})
+        eq(1, meths.get_current_buf().id)
+      end)
+      it('works with :sleep using milliseconds', function()
+        local start = luv.now()
+        meths.cmd({cmd = 'sleep', args = {'100m'}}, {})
+        ok(luv.now() - start <= 300)
+      end)
+    end)
   end)
 end)