Commit 92d54d5e authored by Brett Cannon's avatar Brett Cannon

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

temporarily set or unset environment variables.
parent 5399910e
...@@ -281,4 +281,22 @@ Execute the \class{unittest.TestSuite} instance \var{suite}. ...@@ -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 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 suite so as to print out more detailed information on where the testing suite
originated from. 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} \end{funcdesc}
...@@ -279,7 +279,39 @@ def guard_warnings_filter(): ...@@ -279,7 +279,39 @@ def guard_warnings_filter():
yield yield
finally: finally:
warnings.filters = original_filters 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 # Decorator for running a function in a different locale, correctly resetting
......
...@@ -323,6 +323,10 @@ Extension Modules ...@@ -323,6 +323,10 @@ Extension Modules
Tests 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 - Added guard_warnings_filter to test.test_support. It returns a context
manager that protects the 'warnings' module's filter from being mutated manager that protects the 'warnings' module's filter from being mutated
once the context has been exited. once the context has been exited.
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment