mercurial/changegroup.py

--- a/mercurial/changegroup.py	Wed Feb 20 21:31:42 2008 +0100
+++ b/mercurial/changegroup.py	Thu Feb 21 17:02:28 2008 +0100
@@ -108,22 +108,34 @@
         if cleanup is not None:
             os.unlink(cleanup)
 
-def readbundle(fh, fname):
-    header = fh.read(6)
-    if not header.startswith("HG"):
-        raise util.Abort(_("%s: not a Mercurial bundle file") % fname)
-    elif not header.startswith("HG10"):
-        raise util.Abort(_("%s: unknown bundle version") % fname)
-
-    if header == "HG10BZ":
+def unbundle(header, fh):
+    if header == 'HG10UN':
+        return fh
+    elif not header.startswith('HG'):
+        # old client with uncompressed bundle
+        def generator(f):
+            yield header
+            for chunk in f:
+                yield chunk
+    elif header == 'HG10GZ':
+        def generator(f):
+            zd = zlib.decompressobj()
+            for chunk in f:
+                yield zd.decompress(chunk)
+    elif header == 'HG10BZ':
         def generator(f):
             zd = bz2.BZ2Decompressor()
             zd.decompress("BZ")
             for chunk in util.filechunkiter(f, 4096):
                 yield zd.decompress(chunk)
-        return util.chunkbuffer(generator(fh))
-    elif header == "HG10UN":
-        return fh
+    return util.chunkbuffer(generator(fh))
 
-    raise util.Abort(_("%s: unknown bundle compression type")
-                     % fname)
+def readbundle(fh, fname):
+    header = fh.read(6)
+    if not header.startswith('HG'):
+        raise util.Abort(_('%s: not a Mercurial bundle file') % fname)
+    if not header.startswith('HG10'):
+        raise util.Abort(_('%s: unknown bundle version') % fname)
+    elif header not in bundletypes:
+        raise util.Abort(_('%s: unknown bundle compression type') % fname)
+    return unbundle(header, fh)