Make test.test_support.EnvironmentVarGuard behave like a dictionary.

All changes are mirrored to the underlying os.environ dict, but rolled back
on exit from the with block.

Doc/library/test.rst

@@ -401,9 +401,14 @@ The :mod:`test.test_support` module defines the following classes:
 .. class:: EnvironmentVarGuard()
 
    Class used to temporarily set or unset environment variables.  Instances can be
-   used as a context manager.
+   used as a context manager and have a complete dictionary interface for
+   querying/modifying the underlying ``os.environ``. After exit from the context
+   manager all changes to environment variables done through this instance will
+   be rolled back.
 
    .. versionadded:: 2.6
+   .. versionchanged:: 2.7
+      Added dictionary interface.
 
 
 .. method:: EnvironmentVarGuard.set(envvar, value)
@@ -415,6 +420,7 @@ The :mod:`test.test_support` module defines the following classes:
 
    Temporarily unset the environment variable ``envvar``.
 
+
 .. class:: WarningsRecorder()
 
    Class used to record warnings for unit tests. See documentation of

Lib/test/test_getopt.py

 # test_getopt.py
 # David Goodger <dgoodger@bigfoot.com> 2000-08-19
 
-from test.test_support import verbose, run_doctest, run_unittest
+from test.test_support import verbose, run_doctest, run_unittest, EnvironmentVarGuard
 import unittest
 
 import getopt
@@ -11,15 +11,13 @@ sentinel = object()
 
 class GetoptTests(unittest.TestCase):
     def setUp(self):
-        self.old_posixly_correct = os.environ.get("POSIXLY_CORRECT", sentinel)
-        if self.old_posixly_correct is not sentinel:
-            del os.environ["POSIXLY_CORRECT"]
+        self.env = EnvironmentVarGuard()
+        if "POSIXLY_CORRECT" in self.env:
+            del self.env["POSIXLY_CORRECT"]
 
     def tearDown(self):
-        if self.old_posixly_correct is sentinel:
-            os.environ.pop("POSIXLY_CORRECT", None)
-        else:
-            os.environ["POSIXLY_CORRECT"] = self.old_posixly_correct
+        self.env.__exit__()
+        del self.env
 
     def assertError(self, *args, **kwargs):
         self.assertRaises(getopt.GetoptError, *args, **kwargs)
@@ -135,7 +133,7 @@ class GetoptTests(unittest.TestCase):
         self.assertEqual(args, ['arg1', '-b', '1', '--alpha', '--beta=2'])
 
         # Posix style via POSIXLY_CORRECT
-        os.environ["POSIXLY_CORRECT"] = "1"
+        self.env["POSIXLY_CORRECT"] = "1"
         opts, args = getopt.gnu_getopt(cmdline, 'ab:', ['alpha', 'beta='])
         self.assertEqual(opts, [('-a', '')])
         self.assertEqual(args, ['arg1', '-b', '1', '--alpha', '--beta=2'])

Lib/test/test_gettext.py

@@ -58,10 +58,6 @@ LOCALEDIR = os.path.join('xx', 'LC_MESSAGES')
 MOFILE = os.path.join(LOCALEDIR, 'gettext.mo')
 UMOFILE = os.path.join(LOCALEDIR, 'ugettext.mo')
 MMOFILE = os.path.join(LOCALEDIR, 'metadata.mo')
-try:
-    LANG = os.environ['LANGUAGE']
-except:
-    LANG = 'en'
 
 
 class GettextBaseTest(unittest.TestCase):
@@ -77,10 +73,12 @@ class GettextBaseTest(unittest.TestCase):
         fp = open(MMOFILE, 'wb')
         fp.write(base64.decodestring(MMO_DATA))
         fp.close()
-        os.environ['LANGUAGE'] = 'xx'
+        self.env = test_support.EnvironmentVarGuard()
+        self.env['LANGUAGE'] = 'xx'
 
     def tearDown(self):
-        os.environ['LANGUAGE'] = LANG
+        self.env.__exit__()
+        del self.env
         shutil.rmtree(os.path.split(LOCALEDIR)[0])
 
 

Lib/test/test_ntpath.py

@@ -124,12 +124,11 @@ class TestNtpath(unittest.TestCase):
         tester("ntpath.normpath('//machine/share//a/b')", r'\\machine\share\a\b')
 
     def test_expandvars(self):
-        oldenv = os.environ.copy()
-        try:
-            os.environ.clear()
-            os.environ["foo"] = "bar"
-            os.environ["{foo"] = "baz1"
-            os.environ["{foo}"] = "baz2"
+        with test_support.EnvironmentVarGuard() as env:
+            env.clear()
+            env["foo"] = "bar"
+            env["{foo"] = "baz1"
+            env["{foo}"] = "baz2"
             tester('ntpath.expandvars("foo")', "foo")
             tester('ntpath.expandvars("$foo bar")', "bar bar")
             tester('ntpath.expandvars("${foo}bar")', "barbar")
@@ -149,9 +148,6 @@ class TestNtpath(unittest.TestCase):
             tester('ntpath.expandvars("%?bar%")', "%?bar%")
             tester('ntpath.expandvars("%foo%%bar")', "bar%bar")
             tester('ntpath.expandvars("\'%foo%\'%bar")', "\'%foo%\'%bar")
-        finally:
-            os.environ.clear()
-            os.environ.update(oldenv)
 
     def test_abspath(self):
         # ntpath.abspath() can only be used on a system with the "nt" module

Lib/test/test_optparse.py

@@ -1465,7 +1465,7 @@ class TestHelp(BaseTest):
         # screws things up for other tests when it's part of the Python
         # test suite.
         with test_support.EnvironmentVarGuard() as env:
-            env.set('COLUMNS', str(columns))
+            env['COLUMNS'] = str(columns)
             return InterceptingOptionParser(option_list=options)
 
     def assertHelpEquals(self, expected_output):
@@ -1494,7 +1494,7 @@ class TestHelp(BaseTest):
 
     def test_help_title_formatter(self):
         with test_support.EnvironmentVarGuard() as env:
-            env.set("COLUMNS", "80")
+            env["COLUMNS"] = "80"
             self.parser.formatter = TitledHelpFormatter()
             self.assertHelpEquals(_expected_help_title_formatter)
 

Lib/test/test_posixpath.py

@@ -346,18 +346,17 @@ class PosixPathTest(unittest.TestCase):
             self.assert_(isinstance(posixpath.expanduser("~foo/"), basestring))
 
             with test_support.EnvironmentVarGuard() as env:
-                env.set('HOME', '/')
+                env['HOME'] = '/'
                 self.assertEqual(posixpath.expanduser("~"), "/")
 
         self.assertRaises(TypeError, posixpath.expanduser)
 
     def test_expandvars(self):
-        oldenv = os.environ.copy()
-        try:
-            os.environ.clear()
-            os.environ["foo"] = "bar"
-            os.environ["{foo"] = "baz1"
-            os.environ["{foo}"] = "baz2"
+        with test_support.EnvironmentVarGuard() as env:
+            env.clear()
+            env["foo"] = "bar"
+            env["{foo"] = "baz1"
+            env["{foo}"] = "baz2"
             self.assertEqual(posixpath.expandvars("foo"), "foo")
             self.assertEqual(posixpath.expandvars("$foo bar"), "bar bar")
             self.assertEqual(posixpath.expandvars("${foo}bar"), "barbar")
@@ -370,9 +369,6 @@ class PosixPathTest(unittest.TestCase):
             self.assertEqual(posixpath.expandvars("${{foo}}"), "baz1}")
             self.assertEqual(posixpath.expandvars("$foo$foo"), "barbar")
             self.assertEqual(posixpath.expandvars("$bar$bar"), "$bar$bar")
-        finally:
-            os.environ.clear()
-            os.environ.update(oldenv)
 
         self.assertRaises(TypeError, posixpath.expandvars)
 

Lib/test/test_support.py

@@ -13,6 +13,7 @@ import shutil
 import warnings
 import unittest
 import importlib
+import UserDict
 
 __all__ = ["Error", "TestFailed", "ResourceDenied", "import_module",
            "verbose", "use_resources", "max_memuse", "record_original_stdout",
@@ -526,26 +527,39 @@ class CleanImport(object):
         sys.modules.update(self.original_modules)
 
 
-class EnvironmentVarGuard(object):
+class EnvironmentVarGuard(UserDict.DictMixin):
 
     """Class to help protect the environment variable properly.  Can be used as
     a context manager."""
 
     def __init__(self):
+        self._environ = os.environ
         self._changed = {}
 
-    def set(self, envvar, value):
+    def __getitem__(self, envvar):
+        return self._environ[envvar]
+
+    def __setitem__(self, envvar, value):
         # Remember the initial value on the first access
         if envvar not in self._changed:
-            self._changed[envvar] = os.environ.get(envvar)
-        os.environ[envvar] = value
+            self._changed[envvar] = self._environ.get(envvar)
+        self._environ[envvar] = value
 
-    def unset(self, envvar):
+    def __delitem__(self, envvar):
         # Remember the initial value on the first access
         if envvar not in self._changed:
-            self._changed[envvar] = os.environ.get(envvar)
-        if envvar in os.environ:
-            del os.environ[envvar]
+            self._changed[envvar] = self._environ.get(envvar)
+        if envvar in self._environ:
+            del self._environ[envvar]
+
+    def keys(self):
+        return self._environ.keys()
+
+    def set(self, envvar, value):
+        self[envvar] = value
+
+    def unset(self, envvar):
+        del self[envvar]
 
     def __enter__(self):
         return self
@@ -553,10 +567,10 @@ class EnvironmentVarGuard(object):
     def __exit__(self, *ignore_exc):
         for (k, v) in self._changed.items():
             if v is None:
-                if k in os.environ:
-                    del os.environ[k]
+                if k in self._environ:
+                    del self._environ[k]
             else:
-                os.environ[k] = v
+                self._environ[k] = v
 
 
 class TransientResource(object):

Lib/test/test_tcl.py

@@ -144,23 +144,20 @@ class TclTest(unittest.TestCase):
         import sys
         if sys.platform.startswith(('win', 'darwin', 'cygwin')):
             return  # no failure possible on windows?
-        if 'DISPLAY' in os.environ:
-            old_display = os.environ['DISPLAY']
-            del os.environ['DISPLAY']
-            # on some platforms, deleting environment variables
-            # doesn't actually carry through to the process level
-            # because they don't support unsetenv
-            # If that's the case, abort.
-            display = os.popen('echo $DISPLAY').read().strip()
-            if display:
-                return
-        try:
+        with test_support.EnvironmentVarGuard() as env:
+            if 'DISPLAY' in os.environ:
+                del env['DISPLAY']
+                # on some platforms, deleting environment variables
+                # doesn't actually carry through to the process level
+                # because they don't support unsetenv
+                # If that's the case, abort.
+                display = os.popen('echo $DISPLAY').read().strip()
+                if display:
+                    return
+
             tcl = Tcl()
             self.assertRaises(TclError, tcl.winfo_geometry)
             self.assertRaises(TclError, tcl.loadtk)
-        finally:
-            if old_display is not None:
-                os.environ['DISPLAY'] = old_display
 
 def test_main():
     test_support.run_unittest(TclTest, TkinterTest)

Lib/test/test_tempfile.py

@@ -153,7 +153,7 @@ class test__candidate_tempdir_list(TC):
             for envname in 'TMPDIR', 'TEMP', 'TMP':
                 dirname = os.getenv(envname)
                 if not dirname:
-                    env.set(envname, os.path.abspath(envname))
+                    env[envname] = os.path.abspath(envname)
 
             cand = tempfile._candidate_tempdir_list()
 

Lib/test/test_urllib.py

@@ -103,7 +103,7 @@ class ProxyTests(unittest.TestCase):
         # Delete all proxy related env vars
         for k, v in os.environ.iteritems():
             if 'proxy' in k.lower():
-                env.unset(k)
+                del env[k]
 
     def tearDown(self):
         # Restore all proxy related env vars

Lib/test/test_xmlrpc.py

@@ -605,7 +605,7 @@ class CGIHandlerTestCase(unittest.TestCase):
 
     def test_cgi_get(self):
         with test_support.EnvironmentVarGuard() as env:
-            env.set('REQUEST_METHOD', 'GET')
+            env['REQUEST_METHOD'] = 'GET'
             # if the method is GET and no request_text is given, it runs handle_get
             # get sysout output
             tmp = sys.stdout
@@ -646,7 +646,7 @@ class CGIHandlerTestCase(unittest.TestCase):
         sys.stdout = open(test_support.TESTFN, "w")
 
         with test_support.EnvironmentVarGuard() as env:
-            env.set('CONTENT_LENGTH', str(len(data)))
+            env['CONTENT_LENGTH'] = str(len(data))
             self.cgi.handle_request()
 
         sys.stdin.close()