Merged revisions 72167 via svnmerge from

svn+ssh://pythondev@svn.python.org/python/trunk

........
  r72167 | walter.doerwald | 2009-05-01 19:35:37 +0200 (Fr, 01 Mai 2009) | 5 lines

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

10 files changed, 78 insertions, 67 deletions

diff --git a/Doc/library/test.rst b/Doc/library/test.rst
index d196678..2da63ac 100644
--- a/Doc/library/test.rst
+++ b/Doc/library/test.rst
@@ -384,8 +384,13 @@ The :mod:`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.
 
+   .. versionchanged:: 2.7
+      Added dictionary interface.
 
 .. method:: EnvironmentVarGuard.set(envvar, value)
 
@@ -396,6 +401,7 @@ The :mod:`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
diff --git a/Lib/test/support.py b/Lib/test/support.py
index bdc6164..df114a2 100644
--- a/Lib/test/support.py
+++ b/Lib/test/support.py
@@ -13,6 +13,7 @@ import shutil
 import warnings
 import unittest
 import importlib
+import collections
 
 __all__ = ["Error", "TestFailed", "ResourceDenied", "import_module",
            "verbose", "use_resources", "max_memuse", "record_original_stdout",
@@ -510,26 +511,45 @@ class CleanImport(object):
         sys.modules.update(self.original_modules)
 
 
-class EnvironmentVarGuard(object):
+class EnvironmentVarGuard(collections.MutableMapping):
 
     """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 __iter__(self):
+        return iter(self._environ)
+
+    def __len__(self):
+        return len(self._environ)
+
+    def set(self, envvar, value):
+        self[envvar] = value
+
+    def unset(self, envvar):
+        del self[envvar]
 
     def __enter__(self):
         return self
@@ -537,10 +557,11 @@ 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):
 
diff --git a/Lib/test/test_getopt.py b/Lib/test/test_getopt.py
index 11ae163..b7a6309 100644
--- a/Lib/test/test_getopt.py
+++ b/Lib/test/test_getopt.py
@@ -1,7 +1,7 @@
 # test_getopt.py
 # David Goodger <dgoodger@bigfoot.com> 2000-08-19
 
-from test.support import verbose, run_doctest, run_unittest
+from 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'])
diff --git a/Lib/test/test_gettext.py b/Lib/test/test_gettext.py
index 6ee608e..e26b30a 100644
--- a/Lib/test/test_gettext.py
+++ b/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 = 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])
 
 
diff --git a/Lib/test/test_ntpath.py b/Lib/test/test_ntpath.py
index 3dc7765..2b2b7d5 100644
--- a/Lib/test/test_ntpath.py
+++ b/Lib/test/test_ntpath.py
@@ -141,12 +141,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 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")
@@ -166,9 +165,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
diff --git a/Lib/test/test_optparse.py b/Lib/test/test_optparse.py
index e4ede4a..9c52f86 100644
--- a/Lib/test/test_optparse.py
+++ b/Lib/test/test_optparse.py
@@ -1449,7 +1449,7 @@ class TestHelp(BaseTest):
         # screws things up for other tests when it's part of the Python
         # test suite.
         with support.EnvironmentVarGuard() as env:
-            env.set('COLUMNS', str(columns))
+            env['COLUMNS'] = str(columns)
             return InterceptingOptionParser(option_list=options)
 
     def assertHelpEquals(self, expected_output):
@@ -1474,7 +1474,7 @@ class TestHelp(BaseTest):
 
     def test_help_title_formatter(self):
         with support.EnvironmentVarGuard() as env:
-            env.set("COLUMNS", "80")
+            env["COLUMNS"] = "80"
             self.parser.formatter = TitledHelpFormatter()
             self.assertHelpEquals(_expected_help_title_formatter)
 
diff --git a/Lib/test/test_posixpath.py b/Lib/test/test_posixpath.py
index a37be25..eb918f4 100644
--- a/Lib/test/test_posixpath.py
+++ b/Lib/test/test_posixpath.py
@@ -420,18 +420,17 @@ class PosixPathTest(unittest.TestCase):
             self.assert_(isinstance(posixpath.expanduser(b"~foo/"), bytes))
 
             with 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 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")
@@ -457,11 +456,7 @@ class PosixPathTest(unittest.TestCase):
             self.assertEqual(posixpath.expandvars(b"${{foo}}"), b"baz1}")
             self.assertEqual(posixpath.expandvars(b"$foo$foo"), b"barbar")
             self.assertEqual(posixpath.expandvars(b"$bar$bar"), b"$bar$bar")
-        finally:
-            os.environ.clear()
-            os.environ.update(oldenv)
-
-        self.assertRaises(TypeError, posixpath.expandvars)
+            self.assertRaises(TypeError, posixpath.expandvars)
 
     def test_normpath(self):
         self.assertEqual(posixpath.normpath(""), ".")
diff --git a/Lib/test/test_tcl.py b/Lib/test/test_tcl.py
index f0eb863..cfff971 100644
--- a/Lib/test/test_tcl.py
+++ b/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 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():
     support.run_unittest(TclTest, TkinterTest)
diff --git a/Lib/test/test_tempfile.py b/Lib/test/test_tempfile.py
index 11674a6..69b2e6c 100644
--- a/Lib/test/test_tempfile.py
+++ b/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()
 
diff --git a/Lib/test/test_xmlrpc.py b/Lib/test/test_xmlrpc.py
index 3734750..73d53ee 100644
--- a/Lib/test/test_xmlrpc.py
+++ b/Lib/test/test_xmlrpc.py
@@ -572,7 +572,7 @@ class CGIHandlerTestCase(unittest.TestCase):
 
     def test_cgi_get(self):
         with 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
@@ -613,7 +613,7 @@ class CGIHandlerTestCase(unittest.TestCase):
         sys.stdout = open(support.TESTFN, "w")
 
         with support.EnvironmentVarGuard() as env:
-            env.set('CONTENT_LENGTH', str(len(data)))
+            env['CONTENT_LENGTH'] = str(len(data))
             self.cgi.handle_request()
 
         sys.stdin.close()