gh-132983: Style improvements for compression.zstd (#133547)

Co-authored-by: Adam Turner <9087854+AA-Turner@users.noreply.github.com>
This commit is contained in:
Emma Smith 2025-05-14 05:08:27 -07:00 committed by GitHub
parent e7ad59bd73
commit b44c824856
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 62 additions and 61 deletions

View file

@ -2,28 +2,28 @@
__all__ = (
# compression.zstd
"COMPRESSION_LEVEL_DEFAULT",
"compress",
"CompressionParameter",
"decompress",
"DecompressionParameter",
"finalize_dict",
"get_frame_info",
"Strategy",
"train_dict",
'COMPRESSION_LEVEL_DEFAULT',
'compress',
'CompressionParameter',
'decompress',
'DecompressionParameter',
'finalize_dict',
'get_frame_info',
'Strategy',
'train_dict',
# compression.zstd._zstdfile
"open",
"ZstdFile",
'open',
'ZstdFile',
# _zstd
"get_frame_size",
"zstd_version",
"zstd_version_info",
"ZstdCompressor",
"ZstdDecompressor",
"ZstdDict",
"ZstdError",
'get_frame_size',
'zstd_version',
'zstd_version_info',
'ZstdCompressor',
'ZstdDecompressor',
'ZstdDict',
'ZstdError',
)
import _zstd
@ -43,6 +43,7 @@ COMPRESSION_LEVEL_DEFAULT = _zstd.ZSTD_CLEVEL_DEFAULT
class FrameInfo:
"""Information about a Zstandard frame."""
__slots__ = 'decompressed_size', 'dictionary_id'
def __init__(self, decompressed_size, dictionary_id):
@ -125,13 +126,13 @@ def finalize_dict(zstd_dict, /, samples, dict_size, level):
chunks = b''.join(samples)
chunk_sizes = tuple(_nbytes(sample) for sample in samples)
if not chunks:
raise ValueError("The samples are empty content, can't finalize the"
raise ValueError("The samples are empty content, can't finalize the "
"dictionary.")
dict_content = _zstd.finalize_dict(zstd_dict.dict_content,
chunks, chunk_sizes,
dict_size, level)
dict_content = _zstd.finalize_dict(zstd_dict.dict_content, chunks,
chunk_sizes, dict_size, level)
return ZstdDict(dict_content)
def compress(data, level=None, options=None, zstd_dict=None):
"""Return Zstandard compressed *data* as bytes.
@ -147,6 +148,7 @@ def compress(data, level=None, options=None, zstd_dict=None):
comp = ZstdCompressor(level=level, options=options, zstd_dict=zstd_dict)
return comp.compress(data, mode=ZstdCompressor.FLUSH_FRAME)
def decompress(data, zstd_dict=None, options=None):
"""Decompress one or more frames of Zstandard compressed *data*.
@ -162,12 +164,12 @@ def decompress(data, zstd_dict=None, options=None):
decomp = ZstdDecompressor(options=options, zstd_dict=zstd_dict)
results.append(decomp.decompress(data))
if not decomp.eof:
raise ZstdError("Compressed data ended before the "
"end-of-stream marker was reached")
raise ZstdError('Compressed data ended before the '
'end-of-stream marker was reached')
data = decomp.unused_data
if not data:
break
return b"".join(results)
return b''.join(results)
class CompressionParameter(enum.IntEnum):

View file

@ -4,7 +4,7 @@ from _zstd import (ZstdCompressor, ZstdDecompressor, ZstdError,
ZSTD_DStreamOutSize)
from compression._common import _streams
__all__ = ("ZstdFile", "open")
__all__ = ('ZstdFile', 'open')
_MODE_CLOSED = 0
_MODE_READ = 1
@ -31,15 +31,15 @@ class ZstdFile(_streams.BaseStream):
FLUSH_BLOCK = ZstdCompressor.FLUSH_BLOCK
FLUSH_FRAME = ZstdCompressor.FLUSH_FRAME
def __init__(self, file, /, mode="r", *,
def __init__(self, file, /, mode='r', *,
level=None, options=None, zstd_dict=None):
"""Open a Zstandard compressed file in binary mode.
*file* can be either an file-like object, or a file name to open.
*mode* can be "r" for reading (default), "w" for (over)writing, "x" for
creating exclusively, or "a" for appending. These can equivalently be
given as "rb", "wb", "xb" and "ab" respectively.
*mode* can be 'r' for reading (default), 'w' for (over)writing, 'x' for
creating exclusively, or 'a' for appending. These can equivalently be
given as 'rb', 'wb', 'xb' and 'ab' respectively.
*level* is an optional int specifying the compression level to use,
or COMPRESSION_LEVEL_DEFAULT if not given.
@ -57,33 +57,33 @@ class ZstdFile(_streams.BaseStream):
self._buffer = None
if not isinstance(mode, str):
raise ValueError("mode must be a str")
raise ValueError('mode must be a str')
if options is not None and not isinstance(options, dict):
raise TypeError("options must be a dict or None")
mode = mode.removesuffix("b") # handle rb, wb, xb, ab
if mode == "r":
raise TypeError('options must be a dict or None')
mode = mode.removesuffix('b') # handle rb, wb, xb, ab
if mode == 'r':
if level is not None:
raise TypeError("level is illegal in read mode")
raise TypeError('level is illegal in read mode')
self._mode = _MODE_READ
elif mode in {"w", "a", "x"}:
elif mode in {'w', 'a', 'x'}:
if level is not None and not isinstance(level, int):
raise TypeError("level must be int or None")
raise TypeError('level must be int or None')
self._mode = _MODE_WRITE
self._compressor = ZstdCompressor(level=level, options=options,
zstd_dict=zstd_dict)
self._pos = 0
else:
raise ValueError(f"Invalid mode: {mode!r}")
raise ValueError(f'Invalid mode: {mode!r}')
if isinstance(file, (str, bytes, PathLike)):
self._fp = io.open(file, f'{mode}b')
self._close_fp = True
elif ((mode == 'r' and hasattr(file, "read"))
or (mode != 'r' and hasattr(file, "write"))):
elif ((mode == 'r' and hasattr(file, 'read'))
or (mode != 'r' and hasattr(file, 'write'))):
self._fp = file
else:
raise TypeError("file must be a file-like object "
"or a str, bytes, or PathLike object")
raise TypeError('file must be a file-like object '
'or a str, bytes, or PathLike object')
if self._mode == _MODE_READ:
raw = _streams.DecompressReader(
@ -151,22 +151,22 @@ class ZstdFile(_streams.BaseStream):
return
self._check_not_closed()
if mode not in {self.FLUSH_BLOCK, self.FLUSH_FRAME}:
raise ValueError("Invalid mode argument, expected either "
"ZstdFile.FLUSH_FRAME or "
"ZstdFile.FLUSH_BLOCK")
raise ValueError('Invalid mode argument, expected either '
'ZstdFile.FLUSH_FRAME or '
'ZstdFile.FLUSH_BLOCK')
if self._compressor.last_mode == mode:
return
# Flush zstd block/frame, and write.
data = self._compressor.flush(mode)
self._fp.write(data)
if hasattr(self._fp, "flush"):
if hasattr(self._fp, 'flush'):
self._fp.flush()
def read(self, size=-1):
"""Read up to size uncompressed bytes from the file.
If size is negative or omitted, read until EOF is reached.
Returns b"" if the file is already at EOF.
Returns b'' if the file is already at EOF.
"""
if size is None:
size = -1
@ -178,7 +178,7 @@ class ZstdFile(_streams.BaseStream):
making multiple reads from the underlying stream. Reads up to a
buffer's worth of data if size is negative.
Returns b"" if the file is at EOF.
Returns b'' if the file is at EOF.
"""
self._check_can_read()
if size < 0:
@ -293,7 +293,7 @@ class ZstdFile(_streams.BaseStream):
return self._mode == _MODE_WRITE
def open(file, /, mode="rb", *, level=None, options=None, zstd_dict=None,
def open(file, /, mode='rb', *, level=None, options=None, zstd_dict=None,
encoding=None, errors=None, newline=None):
"""Open a Zstandard compressed file in binary or text mode.
@ -301,8 +301,8 @@ def open(file, /, mode="rb", *, level=None, options=None, zstd_dict=None,
in which case the named file is opened, or it can be an existing file object
to read from or write to.
The mode parameter can be "r", "rb" (default), "w", "wb", "x", "xb", "a",
"ab" for binary mode, or "rt", "wt", "xt", "at" for text mode.
The mode parameter can be 'r', 'rb' (default), 'w', 'wb', 'x', 'xb', 'a',
'ab' for binary mode, or 'rt', 'wt', 'xt', 'at' for text mode.
The level, options, and zstd_dict parameters specify the settings the same
as ZstdFile.
@ -323,19 +323,19 @@ def open(file, /, mode="rb", *, level=None, options=None, zstd_dict=None,
behavior, and line ending(s).
"""
text_mode = "t" in mode
mode = mode.replace("t", "")
text_mode = 't' in mode
mode = mode.replace('t', '')
if text_mode:
if "b" in mode:
raise ValueError(f"Invalid mode: {mode!r}")
if 'b' in mode:
raise ValueError(f'Invalid mode: {mode!r}')
else:
if encoding is not None:
raise ValueError("Argument 'encoding' not supported in binary mode")
raise ValueError('Argument "encoding" not supported in binary mode')
if errors is not None:
raise ValueError("Argument 'errors' not supported in binary mode")
raise ValueError('Argument "errors" not supported in binary mode')
if newline is not None:
raise ValueError("Argument 'newline' not supported in binary mode")
raise ValueError('Argument "newline" not supported in binary mode')
binary_file = ZstdFile(file, mode, level=level, options=options,
zstd_dict=zstd_dict)

View file

@ -2065,7 +2065,7 @@ class TarFile(object):
"gz": "gzopen", # gzip compressed tar
"bz2": "bz2open", # bzip2 compressed tar
"xz": "xzopen", # lzma compressed tar
"zst": "zstopen" # zstd compressed tar
"zst": "zstopen", # zstd compressed tar
}
#--------------------------------------------------------------------------

View file

@ -2512,8 +2512,7 @@ maninstall: altmaninstall
XMLLIBSUBDIRS= xml xml/dom xml/etree xml/parsers xml/sax
LIBSUBDIRS= asyncio \
collections \
compression compression/bz2 compression/gzip compression/zstd \
compression/lzma compression/zlib compression/_common \
compression compression/_common compression/zstd \
concurrent concurrent/futures \
csv \
ctypes ctypes/macholib \