[Python-checkins] r85712 - python/branches/pep-382/Lib/test/test_pep382.py

[eric.smith]

Author: eric.smith
Date: Tue Oct 19 02:49:37 2010
New Revision: 85712

Log:
Changed to use contextlib.contextmanager; refactor the code to find matching module names.

Modified:
   python/branches/pep-382/Lib/test/test_pep382.py

Modified: python/branches/pep-382/Lib/test/test_pep382.py
==============================================================================
--- python/branches/pep-382/Lib/test/test_pep382.py	(original)
+++ python/branches/pep-382/Lib/test/test_pep382.py	Tue Oct 19 02:49:37 2010
@@ -1,4 +1,5 @@
 import unittest, sys, os
+import contextlib
 from test import support
 
 d1="d1"
@@ -7,26 +8,23 @@
 d4="d4"
 test_namespace_prefix = 'pep382test'
 
-class restoring_sys_modules:
-    # in __exit__ removes any modules in sys.modules that start with module_prefix
-    def __init__(self, module_prefix):
-        self.module_prefix = module_prefix
-
-    def __enter__(self):
-        # just make sure no such modules are already present
-        for key in sys.modules:
-            if key.startswith(self.module_prefix):
-                assert False, '{} already in sys.modules'.format(key)
-
-        return self
-
-    def __exit__(self, *ignore_exc):
-        # delete all matching modules
-
-        # convert sys.modules to a list because we're going to mutate it
-        for key in list(sys.modules):
-            if key.startswith(self.module_prefix):
-                del sys.modules[key]
+ at contextlib.contextmanager
+def restoring_sys_modules(module_prefix):
+    # after running, removes any modules in sys.modules that start with module_prefix
+
+    def matching_modules():
+        # return module names that match the prefix
+        return [key for key in sys.modules if key.startswith(module_prefix)]
+
+    # just make sure no matching modules are already present
+    for key in matching_modules():
+        assert False, '{} already in sys.modules'.format(key)
+
+    yield
+
+    # delete all matching modules
+    for key in matching_modules():
+        del sys.modules[key]
 
 
 class PthTestsBase(unittest.TestCase):