Fix typo, handle bzr-specific locking in GenericProcessor.

3 files changed, 27 insertions, 2 deletions

diff --git a/cache_manager.py b/cache_manager.py
index 59f8911..4b0c369 100644
--- a/cache_manager.py
+++ b/cache_manager.py
@@ -79,7 +79,7 @@ class _Cleanup(object):
             self.small_blobs.close()
             self.small_blobs = None
         if self.tempdir is not None:
-            shutils.rmtree(self.tempdir)
+            shutil.rmtree(self.tempdir)
         
 
 class CacheManager(object):
diff --git a/processors/generic_processor.py b/processors/generic_processor.py
index 394dc06..a2fe03c 100644
--- a/processors/generic_processor.py
+++ b/processors/generic_processor.py
@@ -288,6 +288,31 @@ class GenericProcessor(processor.ImportProcessor):
                 self.repo, self.inventory_cache_size,
                 fulltext_when=fulltext_when)
 
+    def process(self, command_iter):
+        """Import data into Bazaar by processing a stream of commands.
+
+        :param command_iter: an iterator providing commands
+        """
+        if self.working_tree is not None:
+            self.working_tree.lock_write()
+        elif self.branch is not None:
+            self.branch.lock_write()
+        elif self.repo is not None:
+            self.repo.lock_write()
+        try:
+            super(GenericProcessor, self)._process(command_iter)
+        finally:
+            # If an unhandled exception occurred, abort the write group
+            if self.repo is not None and self.repo.is_in_write_group():
+                self.repo.abort_write_group()
+            # Release the locks
+            if self.working_tree is not None:
+                self.working_tree.unlock()
+            elif self.branch is not None:
+                self.branch.unlock()
+            elif self.repo is not None:
+                self.repo.unlock()
+
     def _process(self, command_iter):
         # if anything goes wrong, abort the write group if any
         try:
diff --git a/tests/test_filter_processor.py b/tests/test_filter_processor.py
index 2857873..5bc7d97 100644
--- a/tests/test_filter_processor.py
+++ b/tests/test_filter_processor.py
@@ -109,7 +109,7 @@ class TestCaseWithFiltering(TestCase):
     def assertFiltering(self, input, params, expected):
         outf = StringIO()
         proc = filter_processor.FilterProcessor(
-            None, params=params)
+            params=params)
         proc.outf = outf
         s = StringIO(input)
         p = parser.ImportParser(s)