diff --git a/README.md b/README.md index 38b13d7..a60a0ff 100644 --- a/README.md +++ b/README.md @@ -29,9 +29,11 @@ Compression example: ```python import zipfile_zstd as zipfile -zf = zipfile.ZipFile('/tmp/test.zip', 'w', zipfile.ZIP_ZSTANDARD, compresslevel=19) +zf = zipfile.ZipFile('/tmp/test.zip', 'w', zipfile.ZIP_ZSTANDARD, compresslevel=19, threads = -1) zf.write('large_file.img') ``` -Dictionaries and advanced compression parameters are not supported, sorry. +Set `threads=-1` to automatically detect and use all available logical cores, or set the number of desired threads. +The default is to use **half** of the available logical threads/cores. +Dictionaries and other advanced compression parameters are not yet supported. \ No newline at end of file diff --git a/zipfile_zstd/__init__.py b/zipfile_zstd/__init__.py index 61f10fd..07c243d 100644 --- a/zipfile_zstd/__init__.py +++ b/zipfile_zstd/__init__.py @@ -6,5 +6,6 @@ from zipfile import ( ZIP_ZSTANDARD, ZSTANDARD_VERSION, + ZSTANDARD_THREADS ) diff --git a/zipfile_zstd/_zipfile.py b/zipfile_zstd/_zipfile.py index 7896156..33e4d74 100644 --- a/zipfile_zstd/_zipfile.py +++ b/zipfile_zstd/_zipfile.py @@ -3,6 +3,7 @@ import zstandard as zstd import threading import inspect +import os from ._patcher import patch @@ -10,8 +11,15 @@ zipfile.ZIP_ZSTANDARD = 93 zipfile.compressor_names[zipfile.ZIP_ZSTANDARD] = 'zstandard' zipfile.ZSTANDARD_VERSION = 20 - - +zipfile.ZSTANDARD_THREADS = (os.cpu_count()+1)//2 # default = half of the available cores + +@patch(zipfile, 'ZipFile') +def zstd_constructor(*args, **kwargs): + if kwargs and 'threads' in kwargs: + zipfile.ZSTANDARD_THREADS = kwargs.get('threads') + del kwargs['threads'] + return patch.originals['ZipFile'](*args, **kwargs) + @patch(zipfile, '_check_compression') def zstd_check_compression(compression): if compression == zipfile.ZIP_ZSTANDARD: @@ -44,7 +52,7 @@ def zstd_get_compressor(compress_type, compresslevel=None): if compress_type == zipfile.ZIP_ZSTANDARD: if compresslevel is None: compresslevel = 3 - return zstd.ZstdCompressor(level=compresslevel, threads=12).compressobj() + return zstd.ZstdCompressor(level=compresslevel, threads=zipfile.ZSTANDARD_THREADS).compressobj() else: return patch.originals['_get_compressor'](compress_type, compresslevel=compresslevel) else: @@ -53,7 +61,7 @@ def zstd_get_compressor(compress_type, compresslevel=None): if compress_type == zipfile.ZIP_ZSTANDARD: if compresslevel is None: compresslevel = 3 - return zstd.ZstdCompressor(level=compresslevel, threads=12).compressobj() + return zstd.ZstdCompressor(level=compresslevel, threads=zipfile.ZSTANDARD_THREADS).compressobj() else: return patch.originals['_get_compressor'](compress_type)