1 files changed, 33 insertions, 1 deletions
diff --git a/cache_manager.py b/cache_manager.py
index 6c9600f..464403f 100644
--- a/cache_manager.py
+++ b/cache_manager.py
@@ -24,7 +24,35 @@ import time
 import weakref
 
 from bzrlib import lru_cache, trace
-from bzrlib.plugins.fastimport import helpers
+from bzrlib.plugins.fastimport import branch_mapper, helpers
+
+
+class _Cleanup(object):
+    """This class makes sure we clean up when CacheManager goes away.
+
+    We use a helper class to ensure that we are never in a refcycle.
+    """
+
+    def __init__(self, disk_blobs):
+        self.disk_blobs = disk_blobs
+        self.tempdir = None
+        self.small_blobs = None
+
+    def __del__(self):
+        self.finalize()
+
+    def finalize(self):
+        if self.disk_blobs is not None:
+            for info in self.disk_blobs.itervalues():
+                if info[-1] is not None:
+                    os.unlink(info[-1])
+            self.disk_blobs = None
+        if self.small_blobs is not None:
+            self.small_blobs.close()
+            self.small_blobs = None
+        if self.tempdir is not None:
+            shutils.rmtree(self.tempdir)
+        
 
 
 class _Cleanup(object):
@@ -113,6 +141,10 @@ class CacheManager(object):
                 # info not in file - possible when no blobs used
                 pass
 
+        # BranchMapper has no state (for now?), but we keep it around rather
+        # than reinstantiate on every usage
+        self.branch_mapper = branch_mapper.BranchMapper()
+
     def dump_stats(self, note=trace.note):
         """Dump some statistics about what we cached."""
         # TODO: add in inventory stastistics