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

2 files changed, 68 insertions, 4 deletions

diff --git a/runtime/lua/vim/shared.lua b/runtime/lua/vim/shared.lua
index 0c20248d8e..6de1ce9d0c 100644
--- a/runtime/lua/vim/shared.lua
+++ b/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
@@ -1162,18 +1163,19 @@ end
 --- @nodoc
 --- @class vim.context.state
 --- @field bo? table<string, any>
+--- @field env? table<string, any>
 --- @field go? table<string, any>
 --- @field wo? table<string, any>
 
 local scope_map = { buf = 'bo', global = 'go', win = 'wo' }
-local scope_order = { 'o', 'wo', 'bo', 'go' }
-local state_restore_order = { 'bo', 'wo', 'go' }
+local scope_order = { 'o', 'wo', 'bo', 'go', 'env' }
+local state_restore_order = { 'bo', 'wo', 'go', 'env' }
 
 --- Gets data about current state, enough to properly restore specified options/env/etc.
 --- @param context vim.context.mods
 --- @return vim.context.state
 local get_context_state = function(context)
-  local res = { bo = {}, go = {}, wo = {} }
+  local res = { bo = {}, env = {}, go = {}, wo = {} }
 
   -- Use specific order from possibly most to least intrusive
   for _, scope in ipairs(scope_order) do
@@ -1187,7 +1189,9 @@ local get_context_state = function(context)
       -- Always track global option value to properly restore later.
       -- This matters for at least `o` and `wo` (which might set either/both
       -- local and global option values).
-      res.go[name] = res.go[name] or vim.go[name]
+      if sc ~= 'env' then
+        res.go[name] = res.go[name] or vim.go[name]
+      end
     end
   end
 
@@ -1221,6 +1225,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)
diff --git a/test/functional/lua/with_spec.lua b/test/functional/lua/with_spec.lua
index aa2f4bb7f5..99b80ef749 100644
--- a/test/functional/lua/with_spec.lua
+++ b/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 [[
@@ -1546,6 +1604,7 @@ describe('vim._with', function()
     assert_context({ bo = 1 }, 'table')
     assert_context({ buf = 'a' }, 'number')
     assert_context({ emsg_silent = 1 }, 'boolean')
+    assert_context({ env = 1 }, 'table')
     assert_context({ go = 1 }, 'table')
     assert_context({ hide = 1 }, 'boolean')
     assert_context({ keepalt = 1 }, 'boolean')