mercurial/bundlerepo.py

--- a/mercurial/bundlerepo.py	Thu May 11 22:10:47 2006 -0700
+++ b/mercurial/bundlerepo.py	Fri May 12 18:38:56 2006 +0200
@@ -13,7 +13,7 @@
 from node import *
 from i18n import gettext as _
 from demandload import demandload
-demandload(globals(), "changegroup util os struct")
+demandload(globals(), "changegroup util os struct bz2 tempfile")
 
 import localrepo, changelog, manifest, filelog, revlog
 
@@ -159,18 +159,38 @@
 class bundlerepository(localrepo.localrepository):
     def __init__(self, ui, path, bundlename):
         localrepo.localrepository.__init__(self, ui, path)
-        f = open(bundlename, "rb")
-        self.bundlefile = f
+        self.tempfile = None
+        self.bundlefile = open(bundlename, "rb")
         header = self.bundlefile.read(6)
         if not header.startswith("HG"):
             raise util.Abort(_("%s: not a Mercurial bundle file") % bundlename)
         elif not header.startswith("HG10"):
             raise util.Abort(_("%s: unknown bundle version") % bundlename)
         elif header == "HG10BZ":
-            raise util.Abort(_("%s: compressed bundle not supported")
-                             % bundlename)
+            fdtemp, temp = tempfile.mkstemp(prefix="hg-bundle-",
+                                            suffix=".hg10un", dir=self.path)
+            self.tempfile = temp
+            fptemp = os.fdopen(fdtemp, 'wb')
+            def generator(f):
+                zd = bz2.BZ2Decompressor()
+                zd.decompress("BZ")
+                for chunk in f:
+                    yield zd.decompress(chunk)
+            gen = generator(util.filechunkiter(self.bundlefile, 4096))
+
+            try:
+                fptemp.write("HG10UN")
+                for chunk in gen:
+                    fptemp.write(chunk)
+            finally:
+                fptemp.close()
+                self.bundlefile.close()
+
+            self.bundlefile = open(self.tempfile, "rb")
+            # seek right after the header
+            self.bundlefile.seek(6)
         elif header == "HG10UN":
-            # uncompressed bundle supported
+            # nothing to do
             pass
         else:
             raise util.Abort(_("%s: unknown bundle compression type")
@@ -204,3 +224,9 @@
     def close(self):
         """Close assigned bundle file immediately."""
         self.bundlefile.close()
+
+    def __del__(self):
+        if not self.bundlefile.closed:
+            self.bundlefile.close()
+        if self.tempfile is not None:
+            os.unlink(self.tempfile)