Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
May 14, 2025
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.
52 changes: 27 additions & 25 deletions Lib/compression/zstd/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -43,6 +43,7 @@

class FrameInfo:
"""Information about a Zstandard frame."""

__slots__ = 'decompressed_size', 'dictionary_id'

def __init__(self, decompressed_size, dictionary_id):
Expand Down Expand Up @@ -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.

Expand All @@ -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*.

Expand All @@ -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):
Expand Down
66 changes: 33 additions & 33 deletions Lib/compression/zstd/_zstdfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
ZSTD_DStreamOutSize)
from compression._common import _streams

__all__ = ("ZstdFile", "open")
__all__ = ('ZstdFile', 'open')

_MODE_CLOSED = 0
_MODE_READ = 1
Expand All @@ -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.
Expand All @@ -57,33 +57,33 @@ def __init__(self, file, /, mode="r", *,
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(
Expand Down Expand Up @@ -151,22 +151,22 @@ def flush(self, mode=FLUSH_BLOCK):
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
Expand All @@ -178,7 +178,7 @@ def read1(self, size=-1):
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:
Expand Down Expand Up @@ -293,16 +293,16 @@ def writable(self):
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.

file can be either a file name (given as a str, bytes, or PathLike object),
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.
Expand All @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion Lib/tarfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -2065,7 +2065,7 @@ def zstopen(cls, name, mode="r", fileobj=None, level=None, options=None,
"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
}

#--------------------------------------------------------------------------
Expand Down
3 changes: 1 addition & 2 deletions Makefile.pre.in
Original file line number Diff line number Diff line change
Expand Up @@ -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/zstd compression/_common \
concurrent concurrent/futures \
csv \
ctypes ctypes/macholib \
Expand Down
Loading