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

temporarily set or unset environment variables.
This commit is contained in:
Brett Cannon 2007-01-04 00:23:49 +00:00
parent 5399910eba
commit 92d54d5e9c
3 changed files with 55 additions and 1 deletions

View File

@ -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}

View File

@ -280,6 +280,38 @@ def guard_warnings_filter():
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

View File

@ -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.