init

2026-01-17 14:13:43 +09:00
commit ca85a52839
12 changed files with 1692 additions and 0 deletions
--- a/pwr/compression.py
+++ b/pwr/compression.py
@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+import io
+import gzip
+from typing import BinaryIO
+
+from .proto import CompressionAlgorithm, CompressionSettings
+
+
+class CompressionError(Exception):
+    pass
+
+
+def _normalize_algorithm(algorithm: CompressionAlgorithm | int | None) -> CompressionAlgorithm:
+    if algorithm is None:
+        return CompressionAlgorithm.NONE
+    if isinstance(algorithm, CompressionAlgorithm):
+        return algorithm
+    try:
+        return CompressionAlgorithm(int(algorithm))
+    except ValueError as exc:
+        raise CompressionError(f"unknown compression algorithm: {algorithm}") from exc
+
+
+def _brotli_module():
+    try:
+        import brotli  # type: ignore
+
+        return brotli
+    except ImportError:
+        try:
+            import brotlicffi as brotli  # type: ignore
+
+            return brotli
+        except ImportError as exc:
+            raise CompressionError("brotli module not available") from exc
+
+
+class _BrotliReader(io.RawIOBase):
+    def __init__(self, raw: BinaryIO):
+        self._raw = raw
+        brotli = _brotli_module()
+        self._decompressor = brotli.Decompressor()
+        self._buffer = bytearray()
+        self._eof = False
+
+    def readable(self) -> bool:
+        return True
+
+    def _finish(self) -> bytes:
+        if hasattr(self._decompressor, "finish"):
+            return self._decompressor.finish()
+        if hasattr(self._decompressor, "flush"):
+            return self._decompressor.flush()
+        return b""
+
+    def read(self, size: int = -1) -> bytes:
+        if size == 0:
+            return b""
+
+        while not self._eof and (size < 0 or len(self._buffer) < size):
+            chunk = self._raw.read(8192)
+            if not chunk:
+                self._buffer.extend(self._finish())
+                self._eof = True
+                break
+            self._buffer.extend(self._decompressor.process(chunk))
+
+        if size < 0:
+            data = bytes(self._buffer)
+            self._buffer.clear()
+            return data
+
+        data = bytes(self._buffer[:size])
+        del self._buffer[:size]
+        return data
+
+    def readinto(self, b) -> int:
+        data = self.read(len(b))
+        n = len(data)
+        b[:n] = data
+        return n
+
+    def close(self) -> None:
+        if self.closed:
+            return
+        try:
+            if hasattr(self._raw, "close"):
+                self._raw.close()
+        finally:
+            super().close()
+
+
+class _BrotliWriter(io.RawIOBase):
+    def __init__(self, raw: BinaryIO, quality: int | None):
+        self._raw = raw
+        brotli = _brotli_module()
+        kwargs = {}
+        if quality is not None:
+            kwargs["quality"] = int(quality)
+        self._compressor = brotli.Compressor(**kwargs)
+
+    def writable(self) -> bool:
+        return True
+
+    def write(self, b: bytes) -> int:
+        out = self._compressor.process(b)
+        self._raw.write(out)
+        return len(b)
+
+    def close(self) -> None:
+        if self.closed:
+            return
+        try:
+            tail = self._compressor.finish()
+            if tail:
+                self._raw.write(tail)
+            if hasattr(self._raw, "close"):
+                self._raw.close()
+        finally:
+            super().close()
+
+
+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")
+    if algorithm == CompressionAlgorithm.BROTLI:
+        return io.BufferedReader(_BrotliReader(stream))
+    if algorithm == CompressionAlgorithm.ZSTD:
+        try:
+            import zstandard as zstd  # type: ignore
+        except ImportError as exc:
+            raise CompressionError("zstandard module not available") from exc
+        return zstd.ZstdDecompressor().stream_reader(stream)
+    raise CompressionError(f"unsupported compression algorithm: {algorithm}")
+
+
+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)
+    if algorithm == CompressionAlgorithm.BROTLI:
+        return io.BufferedWriter(_BrotliWriter(stream, quality))
+    if algorithm == CompressionAlgorithm.ZSTD:
+        try:
+            import zstandard as zstd  # type: ignore
+        except ImportError as exc:
+            raise CompressionError("zstandard module not available") from exc
+        level = 3 if quality is None else int(quality)
+        return zstd.ZstdCompressor(level=level).stream_writer(stream)
+    raise CompressionError(f"unsupported compression algorithm: {algorithm}")