feat(lua): add context.env (environment variables) to vim._with()

runtime/lua/vim/shared.lua

@@ -1144,6 +1144,7 @@ end
 --- @field bo? table<string, any>
 --- @field buf? integer
 --- @field emsg_silent? boolean
+--- @field env? table<string, any>
 --- @field go? table<string, any>
 --- @field hide? boolean
 --- @field keepalt? boolean
@@ -1159,8 +1160,8 @@ end
 --- @field win? integer
 --- @field wo? table<string, any>
 
-local scope_order = { 'o', 'wo', 'bo', 'go' }
-local cache_restore_order = { 'bo', 'wo', 'go' }
+local scope_order = { 'o', 'wo', 'bo', 'go', 'env' }
+local cache_restore_order = { 'bo', 'wo', 'go', 'env' }
 local scope_map = { buf = 'bo', global = 'go', win = 'wo' }
 
 --- Executes function `f` with the given context specification.
@@ -1190,6 +1191,7 @@ function vim._with(context, f)
   vim.validate('context.bo', context.bo, 'table', true)
   vim.validate('context.buf', context.buf, 'number', true)
   vim.validate('context.emsg_silent', context.emsg_silent, 'boolean', true)
+  vim.validate('context.env', context.env, 'table', true)
   vim.validate('context.go', context.go, 'table', true)
   vim.validate('context.hide', context.hide, 'boolean', true)
   vim.validate('context.keepalt', context.keepalt, 'boolean', true)
@@ -1226,7 +1228,7 @@ function vim._with(context, f)
   -- Decorate so that save-set-restore options is done in correct window-buffer
   local callback = function()
     -- Save and set context (NOTE: triggers `OptionSet` event)
-    local cache = { bo = {}, go = {}, wo = {} } ---@type table<string, any>
+    local cache = { bo = {}, env = {}, go = {}, wo = {} } ---@type table<string, any>
 
     -- Use specific order from possibly most to least intrusive
     for _, scope in ipairs(scope_order) do
@@ -1242,7 +1244,9 @@ function vim._with(context, f)
         -- This matters for at least `o` and `wo` (which might set either/both
         -- local and global option values) and
         -- `wo` (sets both local and global for not global-local window options).
-        cache.go[name] = cache.go[name] or vim.go[name]
+        if sc ~= 'env' then
+          cache.go[name] = cache.go[name] or vim.go[name]
+        end
 
         -- Set using directly original context scope
         vim[scope][name] = context_value

test/functional/lua/with_spec.lua

@@ -338,6 +338,64 @@ describe('vim._with', function()
     end)
   end)
 
+  describe('`env` context', function()
+    before_each(function()
+      exec_lua [[
+        vim.fn.setenv('aaa', 'hello')
+        _G.get_state = function()
+          return { aaa = vim.fn.getenv('aaa'), bbb = vim.fn.getenv('bbb') }
+        end
+      ]]
+    end)
+
+    it('works', function()
+      local out = exec_lua [[
+        local context = { env = { aaa = 'inside', bbb = 'wow' } }
+        local before = get_state()
+        local inner = vim._with(context, get_state)
+        return { before = before, inner = inner, after = get_state() }
+      ]]
+
+      eq({ aaa = 'inside', bbb = 'wow' }, out.inner)
+      eq(out.before, out.after)
+    end)
+
+    it('restores only variables from context', function()
+      local out = exec_lua [[
+        local context = { env = { bbb = 'wow' } }
+        local before = get_state()
+        local inner = vim._with(context, function()
+          vim.env.aaa = 'inside'
+          return get_state()
+        end)
+        return { before = before, inner = inner, after = get_state() }
+      ]]
+
+      eq({ aaa = 'inside', bbb = 'wow' }, out.inner)
+      eq({ aaa = 'inside', bbb = vim.NIL }, out.after)
+    end)
+
+    it('can be nested', function()
+      local out = exec_lua [[
+        local before, before_inner, after_inner = get_state(), nil, nil
+        vim._with({ env = { aaa = 'inside', bbb = 'wow' } }, function()
+          before_inner = get_state()
+          inner = vim._with({ env = { aaa = 'more inside' } }, get_state)
+          after_inner = get_state()
+        end)
+        return {
+          before = before, before_inner = before_inner,
+          inner = inner,
+          after_inner = after_inner, after = get_state(),
+        }
+      ]]
+      eq('more inside', out.inner.aaa)
+      eq('wow', out.inner.bbb)
+      eq(out.before_inner, out.after_inner)
+      eq(out.before, out.after)
+    end)
+  end)
+
   describe('`go` context', function()
     before_each(function()
       exec_lua [[
@@ -1531,6 +1589,7 @@ describe('vim._with', function()
     validate_context({ bo = 1 }, 'table')
     validate_context({ buf = 'a' }, 'number')
     validate_context({ emsg_silent = 1 }, 'boolean')
+    validate_context({ env = 1 }, 'table')
     validate_context({ go = 1 }, 'table')
     validate_context({ hide = 1 }, 'boolean')
     validate_context({ keepalt = 1 }, 'boolean')