Merge pull request #773 from rsyring/gh-769-add-env-swap-kwarg-support

add kwarg support to Env.swap()

docs/tutorial.rst

@@ -134,9 +134,31 @@ variable in Python.  The same is true for deleting them too.
     Become the Lord of the Files
     >>> del $GOAL
 
-Very nice. All environment variables live in the built-in
-``__xonsh_env__`` mapping. You can access this mapping directly, but in most
-situations, you shouldn't need to.
+Very nice.
+
+__xonsh_env__
+--------------
+
+All environment variables live in the built-in ``__xonsh_env__`` mapping. You can access this
+mapping directly, but in most situations, you shouldn't need to.
+
+One helpful method on the __xonsh_env__ is :func:`~xonsh.environ.Env.swap`.  It can be used to temporarily set an
+environment variable:
+
+
+.. code-block:: xonshcon
+
+    >>> with __xonsh_env__.swap(SOMEVAR='foo'):
+    ...     echo $SOMEVAR
+    ...
+    ...
+    foo
+    >>> echo $SOMEVAR
+
+    >>>
+
+Environment Types
+-----------------
 
 Like other variables in Python, environment variables have a type. Sometimes
 this type is imposed based on the variable name. The current rules are pretty

tests/test_environ.py

@@ -72,5 +72,35 @@ def test_HISTCONTROL():
     assert_true('ignoreerr' in env['HISTCONTROL'])
     assert_true('ignoredups' in env['HISTCONTROL'])
 
+def test_swap():
+
+    env = Env(VAR='wakka')
+    assert_equal(env['VAR'], 'wakka')
+
+    # positional arg
+    with env.swap({'VAR': 'foo'}):
+        assert_equal(env['VAR'], 'foo')
+
+    # make sure the environment goes back outside the context manager
+    assert_equal(env['VAR'], 'wakka')
+
+    # kwargs only
+    with env.swap(VAR1='foo', VAR2='bar'):
+        assert_equal(env['VAR1'], 'foo')
+        assert_equal(env['VAR2'], 'bar')
+
+    # positional and kwargs
+    with env.swap({'VAR3': 'baz'}, VAR1='foo', VAR2='bar'):
+        assert_equal(env['VAR1'], 'foo')
+        assert_equal(env['VAR2'], 'bar')
+        assert_equal(env['VAR3'], 'baz')
+
+    # make sure the environment goes back outside the context manager
+    assert_equal(env['VAR'], 'wakka')
+    assert 'VAR1' not in env
+    assert 'VAR2' not in env
+    assert 'VAR3' not in env
+
+
 if __name__ == '__main__':
     nose.runmodule()

xonsh/environ.py

@@ -563,16 +563,27 @@ class Env(MutableMapping):
         return vd
 
     @contextmanager
-    def swap(self, other):
+    def swap(self, other=None, **kwargs):
         """Provides a context manager for temporarily swapping out certain
         environment variables with other values. On exit from the context
         manager, the original values are restored.
         """
         old = {}
-        for k, v in other.items():
+
+        # single positional argument should be a dict-like object
+        if other is not None:
+            for k, v in other.items():
+                old[k] = self.get(k, NotImplemented)
+                self[k] = v
+
+        # kwargs could also have been sent in
+        for k, v in kwargs.items():
             old[k] = self.get(k, NotImplemented)
             self[k] = v
+
         yield self
+
+        # restore the values
         for k, v in old.items():
             if v is NotImplemented:
                 del self[k]