Fixes cleanup code with new namespace management.

enthought · May 1, 2014 · 2de6ed4 · 2de6ed4
1 parent 17f0aed
commit 2de6ed4
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/distarray/cleanup.py b/distarray/cleanup.py
@@ -8,13 +8,15 @@
 
 from distarray.ipython_utils import IPythonClient
 
-def cleanup(view, module_name, context_name):
+def cleanup(view, module_name, prefix):
     """ Delete Context object with the given name from the given module"""
-    def _cleanup(module_name, context_name):
+    def _cleanup(module_name, prefix):
         ns = __import__(module_name)
-        delattr(ns, context_name)
+        for name in vars(ns).copy():
+            if name.startswith(prefix):
+                delattr(ns, name)
 
-    view.apply_async(_cleanup, module_name, context_name)
+    view.apply_sync(_cleanup, module_name, prefix)
 
 
 def cleanup_all(module_name, prefix):

diff --git a/distarray/context.py b/distarray/context.py
@@ -168,7 +168,7 @@ def delete_key(self, key):
 
     def cleanup(self):
         """ Delete keys that this context created from all the engines. """
-        cleanup.cleanup(view=self.view, module_name='__main__', context_name=self.context_key)
+        cleanup.cleanup(view=self.view, module_name='__main__', prefix=self.context_key)
 
     def close(self):
         self.cleanup()