neulus/py-pwr
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix failures

Browse Source
This commit is contained in:
senstella
2026-01-18 03:29:42 +09:00
parent ca85a52839
commit 6560ac93f1
6 changed files with 243 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
pwr/compression.py
View File

@@ -1,7 +1,8 @@
from __future__ import annotations
import io
import gzip
import io
import typing
from typing import BinaryIO
from .proto import CompressionAlgorithm, CompressionSettings
@@ -11,7 +12,9 @@ class CompressionError(Exception):
pass
def _normalize_algorithm(algorithm: CompressionAlgorithm | int | None) -> CompressionAlgorithm:
def _normalize_algorithm(
algorithm: CompressionAlgorithm | int | None,
) -> CompressionAlgorithm:
if algorithm is None:
return CompressionAlgorithm.NONE
if isinstance(algorithm, CompressionAlgorithm):
@@ -103,7 +106,7 @@ class _BrotliWriter(io.RawIOBase):
def writable(self) -> bool:
return True
def write(self, b: bytes) -> int:
def write(self, b) -> int:
out = self._compressor.process(b)
self._raw.write(out)
return len(b)
@@ -121,12 +124,14 @@ class _BrotliWriter(io.RawIOBase):
super().close()
def open_decompressed_reader(stream: BinaryIO, compression: CompressionSettings | None) -> BinaryIO:
def open_decompressed_reader(
stream: BinaryIO, compression: CompressionSettings | None
) -> BinaryIO:
algorithm = _normalize_algorithm(compression.algorithm if compression else None)
if algorithm == CompressionAlgorithm.NONE:
return stream
if algorithm == CompressionAlgorithm.GZIP:
return gzip.GzipFile(fileobj=stream, mode="rb")
return typing.cast(BinaryIO, gzip.GzipFile(fileobj=stream, mode="rb"))
if algorithm == CompressionAlgorithm.BROTLI:
return io.BufferedReader(_BrotliReader(stream))
if algorithm == CompressionAlgorithm.ZSTD:
@@ -138,14 +143,18 @@ def open_decompressed_reader(stream: BinaryIO, compression: CompressionSettings
raise CompressionError(f"unsupported compression algorithm: {algorithm}")
def open_compressed_writer(stream: BinaryIO, compression: CompressionSettings | None) -> BinaryIO:
def open_compressed_writer(
stream: BinaryIO, compression: CompressionSettings | None
) -> BinaryIO:
algorithm = _normalize_algorithm(compression.algorithm if compression else None)
quality = compression.quality if compression else None
if algorithm == CompressionAlgorithm.NONE:
return stream
if algorithm == CompressionAlgorithm.GZIP:
level = 9 if quality is None else int(quality)
return gzip.GzipFile(fileobj=stream, mode="wb", compresslevel=level)
return typing.cast(
BinaryIO, gzip.GzipFile(fileobj=stream, mode="wb", compresslevel=level)
)
if algorithm == CompressionAlgorithm.BROTLI:
return io.BufferedWriter(_BrotliWriter(stream, quality))
if algorithm == CompressionAlgorithm.ZSTD: