Put context's objects in a module instead of a dict.

distarray/cleanup.py

@@ -8,30 +8,13 @@
 
 from distarray.ipython_utils import IPythonClient
 
+def cleanup(view, module_name, context_name):
+    """ Delete Context object with the given name from the given module"""
+    def _cleanup(module_name, context_name):
+        ns = __import__(module_name)
+        delattr(ns, context_name)
 
-def engine_cleanup(module_name, prefix):
-    """ Remove variables with ``prefix`` prefix from the namespace of the
-    module with ``module_name``. 
-    """
-    mod = __import__(module_name)
-    ns = mod.__dict__
-    keys = tuple(ns.keys())
-    deleted = 0
-    for k in keys:
-        if k.startswith(prefix):
-            del ns[k]
-            deleted += 1
-    count = 0
-    for k in ns:
-        if k.startswith(prefix):
-            count += 1
-    return (deleted, count)
-
-
-def cleanup(view, module_name, prefix):
-    """ Delete keys with prefix from client's engines. """
-    remaining = view.apply_async(engine_cleanup, module_name, prefix).get_dict()
-    return remaining
+    view.apply_async(_cleanup, module_name, context_name)
 
 
 def cleanup_all(module_name, prefix):

distarray/context.py

@@ -82,8 +82,10 @@ def _setup_context_key(self):
         Create a dict on the engines which will hold everything from
         this context.
         """
-        context_key = self.uid()
-        cmd = '%s = {}' % (context_key)
+        context_key = DISTARRAY_BASE_NAME + self.uid()
+        cmd = ("import types, sys;"
+               "%s = types.ModuleType('%s');")
+        cmd %= (context_key, context_key)
         self._execute(cmd, targets=range(len(self.view)))
         return context_key
 
@@ -151,7 +153,7 @@ def uid():
 
     def _generate_key(self):
         """ Generate a unique key name for this context. """
-        key = "%s['%s']" % (self.context_key, self.uid())
+        key = "%s.%s" % (self.context_key, 'key_' + self.uid())
         return key
 
     def _key_and_push(self, *values):
@@ -166,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__', prefix=self._key_prefix())
+        cleanup.cleanup(view=self.view, module_name='__main__', context_name=self.context_key)
 
     def close(self):
         self.cleanup()