From 9249a09196b84410fb04153cbe67a8fe22b31929 Mon Sep 17 00:00:00 2001
From: Prodesire <prodesire@yeah.net>
Date: Mon, 15 Jan 2018 00:16:24 +0800
Subject: [PATCH] add environ.environ which is a context manager for updating
 one or more environment variables

---
 pydu/environ.py       | 27 +++++++++++++++++++++++++++
 tests/test_environ.py | 10 ++++++++++
 2 files changed, 37 insertions(+)
 create mode 100644 pydu/environ.py
 create mode 100644 tests/test_environ.py

diff --git a/pydu/environ.py b/pydu/environ.py
new file mode 100644
index 0000000..5d4ec73
--- /dev/null
+++ b/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
diff --git a/tests/test_environ.py b/tests/test_environ.py
new file mode 100644
index 0000000..7ec3ebf
--- /dev/null
+++ b/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