From ca754978001b3aaaf65bddb3493773eb8703f391 Mon Sep 17 00:00:00 2001
From: Kurt Smith <ksmith@enthought.com>
Date: Mon, 21 Apr 2014 14:31:55 -0500
Subject: [PATCH 1/2] Fixes cleanup issues and improves cleanup logic.

---
 distarray/cleanup.py | 38 ++++++++++++++++++++++++++------------
 distarray/context.py | 11 ++++++++---
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/distarray/cleanup.py b/distarray/cleanup.py
index 2adf7394..47242734 100644
--- a/distarray/cleanup.py
+++ b/distarray/cleanup.py
@@ -4,21 +4,34 @@
 #  Distributed under the terms of the BSD License.  See COPYING.rst.
 # ---------------------------------------------------------------------------
 
+from __future__ import print_function
+
 from distarray.ipython_utils import IPythonClient
 
 
-def cleanup(view, prefix):
+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 = ns.keys()
+    deleted = 0
+    for k in keys:
+        if k.startswith(prefix):
+            del ns[k]
+            deleted += 1
+    count = sum([1 for k in ns if k.startswith(prefix)])
+    return (deleted, count)
+
+
+def cleanup(view, module_name, prefix):
     """ Delete keys with prefix from client's engines. """
-    # Delete keys only from this context.
-    def engine_cleanup(prefix):
-        glb = globals()
-        global_keys = list(globals())
-        for gk in global_keys:
-            if gk.startswith(prefix):
-                del glb[gk]
-    view.apply_async(engine_cleanup, prefix)
-
-def cleanup_all(prefix):
+    remaining = view.apply_async(engine_cleanup, module_name, prefix).get_dict()
+    return remaining
+
+
+def cleanup_all(module_name, prefix):
     """ Connects to all engines and runs ``cleanup()`` on them. """
     try:
         c = IPythonClient()
@@ -26,10 +39,11 @@ def cleanup_all(prefix):
         return
     try:
         v = c[:]
-        cleanup(v, prefix)
+        cleanup(v, module_name, prefix)
     finally:
         c.close()
 
+
 def get_local_keys(view, prefix):
     """ Returns a dictionary of keyname -> target_list mapping for all names
         that start with ``prefix`` on engines in ``view``.
diff --git a/distarray/context.py b/distarray/context.py
index 2832242a..8a3c9965 100644
--- a/distarray/context.py
+++ b/distarray/context.py
@@ -19,11 +19,11 @@
 from distarray.externals import six
 from distarray.client import DistArray
 from distarray.client_map import ClientMDMap
-
 from distarray.ipython_utils import IPythonClient
 
+
 DISTARRAY_BASE_NAME = '__distarray__'
-atexit.register(cleanup.cleanup_all, DISTARRAY_BASE_NAME)
+_CLEANUP_REGISTERED = False
 
 
 class Context(object):
@@ -40,6 +40,11 @@ class Context(object):
     '''
 
     def __init__(self, client=None, targets=None):
+        global _CLEANUP_REGISTERED
+
+        if not _CLEANUP_REGISTERED:
+            atexit.register(cleanup.cleanup_all, '__main__', DISTARRAY_BASE_NAME)
+            _CLEANUP_REGISTERED = True
 
         if client is None:
             self.client = IPythonClient()
@@ -162,7 +167,7 @@ def delete_key(self, key):
 
     def cleanup(self):
         """ Delete keys that this context created from all the engines. """
-        cleanup.cleanup(view=self.view, prefix=self._key_prefix())
+        cleanup.cleanup(view=self.view, module_name='__main__', prefix=self._key_prefix())
 
     def close(self):
         self.cleanup()

From 07eeb8ec5d75f7e01e6dd727b7a255d330dbd8ae Mon Sep 17 00:00:00 2001
From: Kurt Smith <ksmith@enthought.com>
Date: Mon, 21 Apr 2014 14:57:41 -0500
Subject: [PATCH 2/2] Fixes py3-related bugs.

---
 distarray/cleanup.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/distarray/cleanup.py b/distarray/cleanup.py
index 47242734..5aa57052 100644
--- a/distarray/cleanup.py
+++ b/distarray/cleanup.py
@@ -15,13 +15,16 @@ def engine_cleanup(module_name, prefix):
     """
     mod = __import__(module_name)
     ns = mod.__dict__
-    keys = ns.keys()
+    keys = tuple(ns.keys())
     deleted = 0
     for k in keys:
         if k.startswith(prefix):
             del ns[k]
             deleted += 1
-    count = sum([1 for k in ns if k.startswith(prefix)])
+    count = 0
+    for k in ns:
+        if k.startswith(prefix):
+            count += 1
     return (deleted, count)