api: make nvim_put support "\022{NUM}" regtype as returned by getregtype()

2 files changed, 47 insertions, 4 deletions

diff --git a/src/nvim/ops.c b/src/nvim/ops.c
index df4452cd4a..bfd02d4ff1 100644
--- a/src/nvim/ops.c
+++ b/src/nvim/ops.c
@@ -5685,9 +5685,6 @@ end:
 /// @param[out] reg Expected to be empty
 bool prepare_yankreg_from_object(yankreg_T *reg, String regtype, size_t lines)
 {
-  if (regtype.size > 1) {
-    return false;
-  }
   char type = regtype.data ? regtype.data[0] : NUL;
 
   switch (type) {
@@ -5707,6 +5704,23 @@ bool prepare_yankreg_from_object(yankreg_T *reg, String regtype, size_t lines)
     return false;
   }
 
+  reg->y_width = 0;
+  if (regtype.size > 1) {
+    if (reg->y_type != kMTBlockWise) {
+      return false;
+    }
+
+    // allow "b7" for a block at least 7 spaces wide
+    if (!ascii_isdigit(regtype.data[1])) {
+      return false;
+    }
+    const char *p = regtype.data+1;
+    reg->y_width = getdigits_int((char_u **)&p)-1;
+    if (regtype.size > (size_t)(p-regtype.data)) {
+      return false;
+    }
+  }
+
   reg->y_array = xcalloc(lines, sizeof(uint8_t *));
   reg->y_size = lines;
   reg->additional_data = NULL;
@@ -5743,7 +5757,7 @@ void finish_yankreg_from_object(yankreg_T *reg, bool clipboard_adjust)
       }
     }
     assert(maxlen <= INT_MAX);
-    reg->y_width = (int)maxlen - 1;
+    reg->y_width = MAX(reg->y_width, (int)maxlen - 1);
   }
 }
 
diff --git a/test/functional/api/vim_spec.lua b/test/functional/api/vim_spec.lua
index 6f7661dd76..02e83c4d58 100644
--- a/test/functional/api/vim_spec.lua
+++ b/test/functional/api/vim_spec.lua
@@ -518,6 +518,35 @@ describe('API', function()
         line 3]])
       eq({0,3,6,0}, funcs.getpos('.'))
     end)
+
+    it('allows block width', function()
+      -- behave consistently with setreg(); support "\022{NUM}" return by getregtype()
+      meths.put({'line 1','line 2','line 3'}, 'l', false, false)
+      expect([[
+        line 1
+        line 2
+        line 3
+        ]])
+
+      -- larger width create spaces
+      meths.put({'a', 'bc'}, 'b3', false, false)
+      expect([[
+        a  line 1
+        bc line 2
+        line 3
+        ]])
+      -- smaller width is ignored
+      meths.put({'xxx', 'yyy'}, '\0221', false, true)
+      expect([[
+        xxxa  line 1
+        yyybc line 2
+        line 3
+        ]])
+      eq({false, "Invalid type: 'bx'"},
+         meth_pcall(meths.put, {'xxx', 'yyy'}, 'bx', false, true))
+      eq({false, "Invalid type: 'b3x'"},
+         meth_pcall(meths.put, {'xxx', 'yyy'}, 'b3x', false, true))
+    end)
   end)
 
   describe('nvim_strwidth', function()