diff --git a/mercurial/util.py b/mercurial/util.py --- a/mercurial/util.py +++ b/mercurial/util.py @@ -3197,5 +3197,50 @@ class _noopengine(compressionengine): compengines.register(_noopengine()) +class _zstdengine(compressionengine): + def name(self): + return 'zstd' + + @propertycache + def _module(self): + # Not all installs have the zstd module available. So defer importing + # until first access. + try: + from . import zstd + # Force delayed import. + zstd.__version__ + return zstd + except ImportError: + return None + + def available(self): + return bool(self._module) + + def bundletype(self): + return 'zstd', 'ZS' + + def compressstream(self, it, opts=None): + opts = opts or {} + # zstd level 3 is almost always significantly faster than zlib + # while providing no worse compression. It strikes a good balance + # between speed and compression. + level = opts.get('level', 3) + + zstd = self._module + z = zstd.ZstdCompressor(level=level).compressobj() + for chunk in it: + data = z.compress(chunk) + if data: + yield data + + yield z.flush() + + def decompressorreader(self, fh): + zstd = self._module + dctx = zstd.ZstdDecompressor() + return chunkbuffer(dctx.read_from(fh)) + +compengines.register(_zstdengine()) + # convenient shortcut dst = debugstacktrace diff --git a/tests/test-bundle-type.t b/tests/test-bundle-type.t --- a/tests/test-bundle-type.t +++ b/tests/test-bundle-type.t @@ -35,17 +35,21 @@ bundle w/o type option test bundle types - $ for t in "None" "bzip2" "gzip" "none-v2" "v2" "v1" "gzip-v1"; do - > echo % test bundle type $t - > hg init t$t + $ testbundle() { + > echo % test bundle type $1 + > hg init t$1 > cd t1 - > hg bundle -t $t ../b$t ../t$t - > f -q -B6 -D ../b$t; echo - > cd ../t$t - > hg debugbundle ../b$t - > hg debugbundle --spec ../b$t + > hg bundle -t $1 ../b$1 ../t$1 + > f -q -B6 -D ../b$1; echo + > cd ../t$1 + > hg debugbundle ../b$1 + > hg debugbundle --spec ../b$1 > echo > cd .. + > } + + $ for t in "None" "bzip2" "gzip" "none-v2" "v2" "v1" "gzip-v1"; do + > testbundle $t > done % test bundle type None searching for changes @@ -106,6 +110,38 @@ test bundle types c35a0f9217e65d1fdb90c936ffa7dbe679f83ddf gzip-v1 +#if zstd + + $ for t in "zstd" "zstd-v2"; do + > testbundle $t + > done + % test bundle type zstd + searching for changes + 1 changesets found + HG20\x00\x00 (esc) + Stream params: sortdict([('Compression', 'ZS')]) + changegroup -- "sortdict([('version', '02'), ('nbchanges', '1')])" + c35a0f9217e65d1fdb90c936ffa7dbe679f83ddf + zstd-v2 + + % test bundle type zstd-v2 + searching for changes + 1 changesets found + HG20\x00\x00 (esc) + Stream params: sortdict([('Compression', 'ZS')]) + changegroup -- "sortdict([('version', '02'), ('nbchanges', '1')])" + c35a0f9217e65d1fdb90c936ffa7dbe679f83ddf + zstd-v2 + +#else + +zstd is a valid engine but isn't available + + $ hg -R t1 bundle -a -t zstd irrelevant.hg + abort: compression engine zstd could not be loaded + [255] + +#endif test garbage file