Support zstd compression via the Python 3.14 compression module

Map --compression values to the compression package submodules so zstd
is available alongside gzip, lzma and bz2. Default remains gzip.
This commit is contained in:
Andrew Hamilton 2026-08-04 21:21:12 +10:00
parent 5e0fd54a92
commit 33b08d3a09
3 changed files with 30 additions and 3 deletions

View file

@ -69,7 +69,7 @@ Options:
-t THEME, --theme=THEME The pygment theme used for syntax
highlighting. Defaults to "native".
-c TYPE, --compression=TYPE The type of compression used in the cache:
gzip, lzma, bz2, or none. Defaults to gzip.
gzip, lzma, bz2, zstd, or none. Defaults to gzip.
--window Use a native GPU window instead of the terminal.
--config=FILE_NAME Use a custom config file instead of tools.toml.
"""
@ -1284,7 +1284,7 @@ def check_arguments():
print("--theme must be one of:", " ".join(themes))
sys.exit(1)
if arguments["--compression"] is not None:
compressions = ["gzip", "lzma", "bz2", "none"]
compressions = ["gzip", "lzma", "bz2", "zstd", "none"]
if arguments["--compression"] not in compressions:
print("--compression must be one of:", " ".join(compressions))
sys.exit(1)

View file

@ -513,9 +513,14 @@ def dump_pickle_safe(object_, path, protocol=pickle.HIGHEST_PROTOCOL, open=open)
os.rename(tmp_path, path)
_COMPRESSION_MODULES = {"gzip": "gzip", "lzma": "lzma", "bz2": "bz2",
"zstd": "zstd"}
@functools.cache
def compression_open_func(compression):
return open if compression == "none" else importlib.import_module(compression).open
return open if compression == "none" else \
importlib.import_module("compression." + _COMPRESSION_MODULES[compression]).open
class Result:

View file

@ -1,7 +1,9 @@
import contextlib
import os
import pickle
import shutil
import tempfile
import unittest
import unittest.mock
@ -158,5 +160,25 @@ class LruCacheWithEvictionTestCase(unittest.TestCase):
self._assert_cache(a, 1, 2, 2)
class CompressionTestCase(unittest.TestCase):
def setUp(self):
self.temp_dir = tempfile.mkdtemp()
def tearDown(self):
shutil.rmtree(self.temp_dir)
def test_open_funcs_round_trip(self):
value = {"text": list(range(1000))}
for compression in ["gzip", "lzma", "bz2", "zstd", "none"]:
with self.subTest(compression=compression):
open_func = tools.compression_open_func(compression)
path = os.path.join(self.temp_dir, f"cache.{compression}")
with open_func(path, "wb") as cache_file:
pickle.dump(value, cache_file)
with open_func(path, "rb") as cache_file:
self.assertEqual(pickle.load(cache_file), value)
if __name__ == "__main__":
golden.main()