diff --git a/distarray/cleanup.py b/distarray/cleanup.py index 333ad2bc..cd18e1e0 100644 --- a/distarray/cleanup.py +++ b/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): diff --git a/distarray/context.py b/distarray/context.py index 489fb2e8..fe8f70e7 100644 --- a/distarray/context.py +++ b/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()