diff options
Diffstat (limited to 'Lib/test/test_support.py')
-rw-r--r-- | Lib/test/test_support.py | 44 |
1 files changed, 44 insertions, 0 deletions
diff --git a/Lib/test/test_support.py b/Lib/test/test_support.py index 2829c5566d..2c19698480 100644 --- a/Lib/test/test_support.py +++ b/Lib/test/test_support.py @@ -3,7 +3,9 @@ if __name__ != 'test.test_support': raise ImportError, 'test_support must be imported from the test package' +from contextlib import contextmanager import sys +import warnings class Error(Exception): """Base class for regression test exceptions.""" @@ -267,6 +269,48 @@ def open_urlresource(url): print >> get_original_stdout(), '\tfetching %s ...' % url fn, _ = urllib.urlretrieve(url, filename) return open(fn) + +@contextmanager +def guard_warnings_filter(): + """Guard the warnings filter from being permanently changed.""" + original_filters = warnings.filters[:] + try: + 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 |