feat(lua): allow vim.wo to be double indexed (#20288)

* feat(lua): allow vim.wo to be double indexed

Problem: `vim.wo` does not implement `setlocal`
Solution: Allow `vim.wo` to be double indexed

Co-authored-by: Christian Clason <c.clason@uni-graz.at>

5 files changed, 65 insertions, 20 deletions

diff --git a/runtime/doc/lua-guide.txt b/runtime/doc/lua-guide.txt
index b138de85c9..28c9ace23b 100644
--- a/runtime/doc/lua-guide.txt
+++ b/runtime/doc/lua-guide.txt
@@ -363,7 +363,7 @@ and `:let &listchars='space:_,tab:>~'`:
 • |vim.o|:  behaves like |:set|
 • |vim.go|: behaves like |:setglobal|
 • |vim.bo|: for buffer-scoped options
-• |vim.wo|: for window-scoped options
+• |vim.wo|: for window-scoped options (can be double indexed)
 
 For example:
 >lua
@@ -386,6 +386,9 @@ window is used:
 >lua
     vim.bo[4].expandtab = true -- sets expandtab to true in buffer 4
     vim.wo.number = true       -- sets number to true in current window
+    vim.wo[0].number = true    -- same as above 
+    vim.wo[0][0].number = true -- sets number to true in current buffer
+                               -- in current window only 
     print(vim.wo[0].number)    --> true
 <
 ------------------------------------------------------------------------------
diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt
index ca4adbce9c..63523c32d5 100644
--- a/runtime/doc/lua.txt
+++ b/runtime/doc/lua.txt
@@ -1175,20 +1175,21 @@ vim.bo[{bufnr}]                                                                *
         print(vim.bo.comments)
         print(vim.bo.baz)                 -- error: invalid key
 <
-vim.wo[{winid}]                                                                *vim.wo*
-    Get or set window-scoped |options| for the window with handle {winid}.
-    Like `:set`. If [{winid}] is omitted then the current window is used.
-    Invalid {winid} or key is an error.
+vim.wo[{winid}][{bufnr}]                                                       *vim.wo*
+    Get or set window-scoped |options| for the window with handle {winid} and
+    buffer with number {bufnr}. Like `:setlocal` if {bufnr} is provided, like
+    `:set` otherwise. If [{winid}] is omitted then the current window is
+    used. Invalid {winid}, {bufnr} or key is an error.
+
+    Note: only {bufnr} with value `0` (the current buffer in the window) is
+    supported.
 
-    Note: this does not access |local-options| (`:setlocal`) instead use: >lua
-        nvim_get_option_value(OPTION, { scope = 'local', win = winid })
-        nvim_set_option_value(OPTION, VALUE, { scope = 'local', win = winid }
-<
     Example: >lua
         local winid = vim.api.nvim_get_current_win()
         vim.wo[winid].number = true    -- same as vim.wo.number = true
         print(vim.wo.foldmarker)
         print(vim.wo.quux)             -- error: invalid key
+        vim.wo[winid][0].spell = false -- like ':setlocal nospell'
 <
 
 
diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 24e9dc917b..ac0a57e4d3 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -110,6 +110,9 @@ The following new APIs and features were added.
 • |vim.ui.open()| opens URIs using the system default handler (macOS `open`,
   Windows `explorer`, Linux `xdg-open`, etc.)
 
+• |vim.wo| can now be double indexed for |:setlocal| behaviour. Currently only
+  `0` for the buffer index is currently supported.
+
 ==============================================================================
 CHANGED FEATURES                                                 *news-changed*
 
diff --git a/runtime/lua/vim/_meta.lua b/runtime/lua/vim/_meta.lua
index 913f1fe203..41e6e8be86 100644
--- a/runtime/lua/vim/_meta.lua
+++ b/runtime/lua/vim/_meta.lua
@@ -67,25 +67,60 @@ local function opt_validate(option_name, target_scope)
   end
 end
 
-local function new_opt_accessor(handle, scope)
+local function new_buf_opt_accessor(bufnr)
   return setmetatable({}, {
     __index = function(_, k)
-      if handle == nil and type(k) == 'number' then
-        return new_opt_accessor(k, scope)
+      if bufnr == nil and type(k) == 'number' then
+        return new_buf_opt_accessor(k)
       end
-      opt_validate(k, scope)
-      return api.nvim_get_option_value(k, { [scope] = handle or 0 })
+      opt_validate(k, 'buf')
+      return api.nvim_get_option_value(k, { buf = bufnr or 0 })
     end,
 
     __newindex = function(_, k, v)
-      opt_validate(k, scope)
-      return api.nvim_set_option_value(k, v, { [scope] = handle or 0 })
+      opt_validate(k, 'buf')
+      return api.nvim_set_option_value(k, v, { buf = bufnr or 0 })
     end,
   })
 end
 
-vim.bo = new_opt_accessor(nil, 'buf')
-vim.wo = new_opt_accessor(nil, 'win')
+vim.bo = new_buf_opt_accessor()
+
+local function new_win_opt_accessor(winid, bufnr)
+  return setmetatable({}, {
+    __index = function(_, k)
+      if bufnr == nil and type(k) == 'number' then
+        if winid == nil then
+          return new_win_opt_accessor(k)
+        else
+          return new_win_opt_accessor(winid, k)
+        end
+      end
+
+      if bufnr ~= nil and bufnr ~= 0 then
+        error('only bufnr=0 is supported')
+      end
+
+      opt_validate(k, 'win')
+      -- TODO(lewis6991): allow passing both buf and win to nvim_get_option_value
+      return api.nvim_get_option_value(k, {
+        scope = bufnr and 'local' or nil,
+        win = winid or 0,
+      })
+    end,
+
+    __newindex = function(_, k, v)
+      opt_validate(k, 'win')
+      -- TODO(lewis6991): allow passing both buf and win to nvim_set_option_value
+      return api.nvim_set_option_value(k, v, {
+        scope = bufnr and 'local' or nil,
+        win = winid or 0,
+      })
+    end,
+  })
+end
+
+vim.wo = new_win_opt_accessor()
 
 -- vim global option
 --  this ONLY sets the global option. like `setglobal`
diff --git a/test/functional/lua/vim_spec.lua b/test/functional/lua/vim_spec.lua
index d5f550a5d1..f168e6ba1d 100644
--- a/test/functional/lua/vim_spec.lua
+++ b/test/functional/lua/vim_spec.lua
@@ -1532,8 +1532,6 @@ describe('lua stdlib', function()
     eq(0, funcs.luaeval "vim.wo[1001].cole")
     matches("Invalid option %(not found%): 'notanopt'$",
        pcall_err(exec_lua, 'return vim.wo.notanopt'))
-    matches("Expected lua string$",
-       pcall_err(exec_lua, 'return vim.wo[0][0].list'))
     matches("Invalid window id: %-1$",
        pcall_err(exec_lua, 'return vim.wo[-1].list'))
     eq(2, funcs.luaeval "vim.wo[1000].cole")
@@ -1548,6 +1546,11 @@ describe('lua stdlib', function()
     eq(200, funcs.luaeval "vim.wo.scrolloff")
     exec_lua [[vim.wo.scrolloff = -1]]
     eq(100, funcs.luaeval "vim.wo.scrolloff")
+    exec_lua [[
+    vim.wo[0][0].scrolloff = 200
+    vim.cmd "split"
+    ]]
+    eq(100, funcs.luaeval "vim.wo.scrolloff")
   end)
 
   describe('vim.opt', function()