Add EnvironmentVarGuard to test.test_support. Provides a context manager to

temporarily set or unset environment variables.

diff --git a/Doc/lib/libtest.tex b/Doc/lib/libtest.tex
index 811d2f7fc08958abc0695abfa095ec0a550936ea..f89c70723c94d8525e4a30c982d22d15075b939b 100644 (file)
--- a/Doc/lib/libtest.tex
+++ b/Doc/lib/libtest.tex
@@ -281,4 +281,22 @@ Execute the \class{unittest.TestSuite} instance \var{suite}.
 The optional argument \var{testclass} accepts one of the test classes in the
 suite so as to print out more detailed information on where the testing suite
 originated from.
+
+The \module{test.test_support} module defines the following classes:
+
+\begin{classdesc}{EnvironmentVarGuard}{}
+Class used to temporarily set or unset environment variables.  Instances can be
+used as a context manager.
+\versionadded{2.6}
+\end{classdesc}
+
+\begin{methoddesc}{set}{envvar, value}
+Temporarily set the environment variable \code{envvar} to the value of
+\code{value}.
+\end{methoddesc}
+
+\begin{methoddesc}{unset}{envvar}
+Temporarily unset the environment variable \code{envvar}.
+\end{methoddesc}
+
 \end{funcdesc}

diff --git a/Lib/test/test_support.py b/Lib/test/test_support.py
index 493941097999322b2481a146da549ef74cd9034d..fa02eac613c9b9903052dcfa3cd73962dc70f75b 100644 (file)
--- a/Lib/test/test_support.py
+++ b/Lib/test/test_support.py
@@ -279,7 +279,39 @@ def guard_warnings_filter():
         yield
     finally:
         warnings.filters = original_filters
-    
+
+class EnvironmentVarGuard(object):
+
+    """Class to help protect the environment variable properly.  Can be used as
+    a context manager."""
+
+    def __init__(self):
+        from os import environ
+        self._environ = environ
+        self._unset = set()
+        self._reset = dict()
+
+    def set(self, envvar, value):
+        if envvar not in self._environ:
+            self._unset.add(envvar)
+        else:
+            self._reset[envvar] = self._environ[envvar]
+        self._environ[envvar] = value
+
+    def unset(self, envvar):
+        if envvar in self._environ:
+            self._reset[envvar] = self._environ[envvar]
+            del self._environ[envvar]
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *ignore_exc):
+        for envvar, value in self._reset.iteritems():
+            self._environ[envvar] = value
+        for unset in self._unset:
+            del self._environ[unset]
+
 
 #=======================================================================
 # Decorator for running a function in a different locale, correctly resetting

diff --git a/Misc/NEWS b/Misc/NEWS
index ba882ef89d90d12419bf3495cb22d4692d247719..4f310d8796ee4ed1a44700b39d796ca0c7a670d4 100644 (file)
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -323,6 +323,10 @@ Extension Modules
 Tests
 -----
 
+- Added test.test_support.EnvironmentVarGuard.  It's a class that provides a
+  context manager so that one can temporarily set or unset environment
+  variables.
+
 - Added guard_warnings_filter to test.test_support.  It returns a context
   manager that protects the 'warnings' module's filter from being mutated
   once the context has been exited.