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 811d2f7..f89c707 100644
--- a/Doc/lib/libtest.tex
+++ b/Doc/lib/libtest.tex
@@ -281,4 +281,22 @@
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 4939410..fa02eac 100644
--- a/Lib/test/test_support.py
+++ b/Lib/test/test_support.py
@@ -279,7 +279,39 @@
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 ba882ef..4f310d8 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -323,6 +323,10 @@
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.