diff --git a/Doc/lib/libtest.tex b/Doc/lib/libtest.tex index 811d2f7fc08..f89c70723c9 100644 --- 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 49394109799..fa02eac613c 100644 --- 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 ba882ef89d9..4f310d8796e 100644 --- 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.