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.

1 files changed, 25 insertions, 11 deletions

diff --git a/Lib/test/test_support.py b/Lib/test/test_support.py
index c1b1c69..7d0ec8d 100644
--- a/Lib/test/test_support.py
+++ b/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):