diff --git a/mercurial/changegroup.py b/mercurial/changegroup.py --- a/mercurial/changegroup.py +++ b/mercurial/changegroup.py @@ -8,7 +8,7 @@ of the GNU General Public License, incor """ from i18n import gettext as _ from demandload import * -demandload(globals(), "struct util") +demandload(globals(), "struct os bz2 util tempfile") def getchunk(source): """get a chunk from a changegroup""" @@ -41,3 +41,55 @@ def genchunk(data): def closechunk(): return struct.pack(">l", 0) +class nocompress(object): + def compress(self, x): + return x + def flush(self): + return "" + +def writebundle(cg, filename, compress): + """Write a bundle file and return its filename. + + Existing files will not be overwritten. + If no filename is specified, a temporary file is created. + bz2 compression can be turned off. + The bundle file will be deleted in case of errors. + """ + + fh = None + cleanup = None + try: + if filename: + if os.path.exists(filename): + raise util.Abort(_("file '%s' already exists") % filename) + fh = open(filename, "wb") + else: + fd, filename = tempfile.mkstemp(prefix="hg-bundle-", suffix=".hg") + fh = os.fdopen(fd, "wb") + cleanup = filename + + if compress: + fh.write("HG10") + z = bz2.BZ2Compressor(9) + else: + fh.write("HG10UN") + z = nocompress() + # parse the changegroup data, otherwise we will block + # in case of sshrepo because we don't know the end of the stream + + # an empty chunkiter is the end of the changegroup + empty = False + while not empty: + empty = True + for chunk in chunkiter(cg): + empty = False + fh.write(z.compress(genchunk(chunk))) + fh.write(z.compress(closechunk())) + fh.write(z.flush()) + cleanup = None + return filename + finally: + if fh is not None: + fh.close() + if cleanup is not None: + os.unlink(cleanup)