From 2de6ed484db3eff580503185021d0be8e539e2f6 Mon Sep 17 00:00:00 2001 From: Kurt Smith Date: Wed, 30 Apr 2014 19:54:13 -0500 Subject: [PATCH] Fixes cleanup code with new namespace management. --- distarray/cleanup.py | 10 ++++++---- distarray/context.py | 2 +- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/distarray/cleanup.py b/distarray/cleanup.py index cd18e1e0..06d444ab 100644 --- 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 index 41b6656d..270ba293 100644 --- 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()