add environ.environ which is a context manager for updating one or more environment variables

pydu/environ.py

@@ -0,0 +1,27 @@
+import os
+from pydu.compat import iteritems
+from contextlib import contextmanager
+
+
+@contextmanager
+def environ(**kwargs):
+    """
+    Context manager for updating one or more environment variables.
+
+    Preserves the previous environment variable (if available) and
+    recovers when exiting the context manager.
+    """
+    original_kwargs = {}
+
+    for key in kwargs:
+        original_kwargs[key] = os.environ.get(key, None)
+        os.environ[key] = kwargs[key]
+
+    yield
+
+    for key, value in iteritems(original_kwargs):
+        if value is None:
+            del os.environ[key]
+            continue
+
+        os.environ[key] = value

tests/test_environ.py

@@ -0,0 +1,10 @@
+import os
+from pydu.environ import environ
+
+
+def test_environ():
+    with environ(a='a', b=''):
+        assert os.environ['a'] == 'a'
+        assert os.environ['b'] == ''
+    assert 'a' not in os.environ
+    assert 'b' not in os.environ