refactor: Extract zlib-related logic into a single module (#7223) (#7250

)

<!-- Thank you for your contribution! -->

Backport #7223 

Addresses issue #7192
Refactors the logic to have the zlib-related stuff concentrated into a
single module

No

#7192

- [x] I think the code is well written
- [x] Unit tests for the changes exist
- [ ] Documentation reflects the changes
- [ ] If you provide code modification, please add yourself to
`CONTRIBUTORS.txt`
  * The format is &lt;Name&gt; &lt;Surname&gt;.
  * Please keep alphabetical order, the file is sorted by names.
- [ ] Add a new news fragment into the `CHANGES` folder
  * name it `<issue_id>.<type>` for example (588.bugfix)
* if you don't have an `issue_id` change it to the pr id after creating
the pr
* ensure type is one of the following: * `.feature`: Signifying a new
feature. * `.bugfix`: Signifying a bug fix. * `.doc`: Signifying a
documentation improvement. * `.removal`: Signifying a deprecation or
removal of public API.
* `.misc`: A ticket has been closed, but it is not of interest to users.
* Make sure to use full sentences with correct case and punctuation, for
example: "Fix issue with non-ascii contents in doctest text files."

---------

Co-authored-by: pre-commit-ci[bot]
<66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Sam Bull <[email protected]>
(cherry picked from commit 3ff81dc)

<!-- Thank you for your contribution! -->

## What do these changes do?

<!-- Please give a short brief about these changes. -->

## Are there changes in behavior for the user?

<!-- Outline any notable behaviour for the end users. -->

## Related issue number

<!-- Are there any issues opened that will be resolved by merging this
change? -->

## Checklist

- [ ] I think the code is well written
- [ ] Unit tests for the changes exist
- [ ] Documentation reflects the changes
- [ ] If you provide code modification, please add yourself to
`CONTRIBUTORS.txt`
  * The format is &lt;Name&gt; &lt;Surname&gt;.
  * Please keep alphabetical order, the file is sorted by names.
- [ ] Add a new news fragment into the `CHANGES` folder
  * name it `<issue_id>.<type>` for example (588.bugfix)
* if you don't have an `issue_id` change it to the pr id after creating
the pr
  * ensure type is one of the following:
    * `.feature`: Signifying a new feature.
    * `.bugfix`: Signifying a bug fix.
    * `.doc`: Signifying a documentation improvement.
    * `.removal`: Signifying a deprecation or removal of public API.
* `.misc`: A ticket has been closed, but it is not of interest to users.
* Make sure to use full sentences with correct case and punctuation, for
example: "Fix issue with non-ascii contents in doctest text files."

---------

Co-authored-by: Mykola Mokhnach <[email protected]>

aiohttp/client_reqrep.py

@@ -38,6 +38,7 @@
     InvalidURL,
     ServerFingerprintMismatch,
 )
+from .compression_utils import HAS_BROTLI
 from .formdata import FormData
 from .helpers import (
     BaseTimerContext,
@@ -51,7 +52,6 @@
     set_result,
 )
 from .http import SERVER_SOFTWARE, HttpVersion10, HttpVersion11, StreamWriter
-from .http_parser import HAS_BROTLI
 from .log import client_logger
 from .streams import StreamReader
 from .typedefs import (

aiohttp/compression_utils.py

@@ -0,0 +1,148 @@
+import asyncio
+import zlib
+from concurrent.futures import Executor
+from typing import Optional, cast
+
+try:
+    import brotli
+
+    HAS_BROTLI = True
+except ImportError:  # pragma: no cover
+    HAS_BROTLI = False
+
+MAX_SYNC_CHUNK_SIZE = 1024
+
+
+def encoding_to_mode(
+    encoding: Optional[str] = None,
+    suppress_deflate_header: bool = False,
+) -> int:
+    if encoding == "gzip":
+        return 16 + zlib.MAX_WBITS
+
+    return -zlib.MAX_WBITS if suppress_deflate_header else zlib.MAX_WBITS
+
+
+class ZlibBaseHandler:
+    def __init__(
+        self,
+        mode: int,
+        executor: Optional[Executor] = None,
+        max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE,
+    ):
+        self._mode = mode
+        self._executor = executor
+        self._max_sync_chunk_size = max_sync_chunk_size
+
+
+class ZLibCompressor(ZlibBaseHandler):
+    def __init__(
+        self,
+        encoding: Optional[str] = None,
+        suppress_deflate_header: bool = False,
+        level: Optional[int] = None,
+        wbits: Optional[int] = None,
+        strategy: int = zlib.Z_DEFAULT_STRATEGY,
+        executor: Optional[Executor] = None,
+        max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE,
+    ):
+        super().__init__(
+            mode=encoding_to_mode(encoding, suppress_deflate_header)
+            if wbits is None
+            else wbits,
+            executor=executor,
+            max_sync_chunk_size=max_sync_chunk_size,
+        )
+        if level is None:
+            self._compressor = zlib.compressobj(wbits=self._mode, strategy=strategy)
+        else:
+            self._compressor = zlib.compressobj(
+                wbits=self._mode, strategy=strategy, level=level
+            )
+
+    def compress_sync(self, data: bytes) -> bytes:
+        return self._compressor.compress(data)
+
+    async def compress(self, data: bytes) -> bytes:
+        if (
+            self._max_sync_chunk_size is not None
+            and len(data) > self._max_sync_chunk_size
+        ):
+            return await asyncio.get_event_loop().run_in_executor(
+                self._executor, self.compress_sync, data
+            )
+        return self.compress_sync(data)
+
+    def flush(self, mode: int = zlib.Z_FINISH) -> bytes:
+        return self._compressor.flush(mode)
+
+
+class ZLibDecompressor(ZlibBaseHandler):
+    def __init__(
+        self,
+        encoding: Optional[str] = None,
+        suppress_deflate_header: bool = False,
+        executor: Optional[Executor] = None,
+        max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE,
+    ):
+        super().__init__(
+            mode=encoding_to_mode(encoding, suppress_deflate_header),
+            executor=executor,
+            max_sync_chunk_size=max_sync_chunk_size,
+        )
+        self._decompressor = zlib.decompressobj(wbits=self._mode)
+
+    def decompress_sync(self, data: bytes, max_length: int = 0) -> bytes:
+        return self._decompressor.decompress(data, max_length)
+
+    async def decompress(self, data: bytes, max_length: int = 0) -> bytes:
+        if (
+            self._max_sync_chunk_size is not None
+            and len(data) > self._max_sync_chunk_size
+        ):
+            return await asyncio.get_event_loop().run_in_executor(
+                self._executor, self.decompress_sync, data, max_length
+            )
+        return self.decompress_sync(data, max_length)
+
+    def flush(self, length: int = 0) -> bytes:
+        return (
+            self._decompressor.flush(length)
+            if length > 0
+            else self._decompressor.flush()
+        )
+
+    @property
+    def eof(self) -> bool:
+        return self._decompressor.eof
+
+    @property
+    def unconsumed_tail(self) -> bytes:
+        return self._decompressor.unconsumed_tail
+
+    @property
+    def unused_data(self) -> bytes:
+        return self._decompressor.unused_data
+
+
+class BrotliDecompressor:
+    # Supports both 'brotlipy' and 'Brotli' packages
+    # since they share an import name. The top branches
+    # are for 'brotlipy' and bottom branches for 'Brotli'
+    def __init__(self) -> None:
+        if not HAS_BROTLI:
+            raise RuntimeError(
+                "The brotli decompression is not available. "
+                "Please install `Brotli` module"
+            )
+        self._obj = brotli.Decompressor()
+
+    def decompress_sync(self, data: bytes) -> bytes:
+        if hasattr(self._obj, "decompress"):
+            return cast(bytes, self._obj.decompress(data))
+        return cast(bytes, self._obj.process(data))
+
+    def flush(self) -> bytes:
+        if hasattr(self._obj, "flush"):
+            return cast(bytes, self._obj.flush())
+        return b""

aiohttp/http_parser.py

@@ -3,7 +3,6 @@
 import collections
 import re
 import string
-import zlib
 from contextlib import suppress
 from enum import IntEnum
 from typing import (
@@ -18,14 +17,14 @@
     Type,
     TypeVar,
     Union,
-    cast,
 )
 
 from multidict import CIMultiDict, CIMultiDictProxy, istr
 from yarl import URL
 
 from . import hdrs
 from .base_protocol import BaseProtocol
+from .compression_utils import HAS_BROTLI, BrotliDecompressor, ZLibDecompressor
 from .helpers import NO_EXTENSIONS, BaseTimerContext
 from .http_exceptions import (
     BadHttpMessage,
@@ -41,14 +40,6 @@
 from .streams import EMPTY_PAYLOAD, StreamReader
 from .typedefs import Final, RawHeaders
 
-try:
-    import brotli
-
-    HAS_BROTLI = True
-except ImportError:  # pragma: no cover
-    HAS_BROTLI = False
-
-
 __all__ = (
     "HeadersParser",
     "HttpParser",
@@ -868,34 +859,16 @@ def __init__(self, out: StreamReader, encoding: Optional[str]) -> None:
         self.encoding = encoding
         self._started_decoding = False
 
+        self.decompressor: Union[BrotliDecompressor, ZLibDecompressor]
         if encoding == "br":
             if not HAS_BROTLI:  # pragma: no cover
                 raise ContentEncodingError(
                     "Can not decode content-encoding: brotli (br). "
                     "Please install `Brotli`"
                 )
-
-            class BrotliDecoder:
-                # Supports both 'brotlipy' and 'Brotli' packages
-                # since they share an import name. The top branches
-                # are for 'brotlipy' and bottom branches for 'Brotli'
-                def __init__(self) -> None:
-                    self._obj = brotli.Decompressor()
-
-                def decompress(self, data: bytes) -> bytes:
-                    if hasattr(self._obj, "decompress"):
-                        return cast(bytes, self._obj.decompress(data))
-                    return cast(bytes, self._obj.process(data))
-
-                def flush(self) -> bytes:
-                    if hasattr(self._obj, "flush"):
-                        return cast(bytes, self._obj.flush())
-                    return b""
-
-            self.decompressor = BrotliDecoder()
+            self.decompressor = BrotliDecompressor()
         else:
-            zlib_mode = 16 + zlib.MAX_WBITS if encoding == "gzip" else zlib.MAX_WBITS
-            self.decompressor = zlib.decompressobj(wbits=zlib_mode)
+            self.decompressor = ZLibDecompressor(encoding=encoding)
 
     def set_exception(self, exc: BaseException) -> None:
         self.out.set_exception(exc)
@@ -916,10 +889,12 @@ def feed_data(self, chunk: bytes, size: int) -> None:
         ):
             # Change the decoder to decompress incorrectly compressed data
             # Actually we should issue a warning about non-RFC-compliant data.
-            self.decompressor = zlib.decompressobj(wbits=-zlib.MAX_WBITS)
+            self.decompressor = ZLibDecompressor(
+                encoding=self.encoding, suppress_deflate_header=True
+            )
 
         try:
-            chunk = self.decompressor.decompress(chunk)
+            chunk = self.decompressor.decompress_sync(chunk)
         except Exception:
             raise ContentEncodingError(
                 "Can not decode content-encoding: %s" % self.encoding

aiohttp/http_websocket.py

@@ -13,6 +13,7 @@
 from typing import Any, Callable, List, Optional, Pattern, Set, Tuple, Union, cast
 
 from .base_protocol import BaseProtocol
+from .compression_utils import ZLibCompressor, ZLibDecompressor
 from .helpers import NO_EXTENSIONS
 from .streams import DataQueue
 from .typedefs import Final
@@ -278,7 +279,7 @@ def __init__(
         self._payload_length = 0
         self._payload_length_flag = 0
         self._compressed: Optional[bool] = None
-        self._decompressobj: Any = None  # zlib.decompressobj actually
+        self._decompressobj: Optional[ZLibDecompressor] = None
         self._compress = compress
 
     def feed_eof(self) -> None:
@@ -298,7 +299,7 @@ def feed_data(self, data: bytes) -> Tuple[bool, bytes]:
     def _feed_data(self, data: bytes) -> Tuple[bool, bytes]:
         for fin, opcode, payload, compressed in self.parse_frame(data):
             if compressed and not self._decompressobj:
-                self._decompressobj = zlib.decompressobj(wbits=-zlib.MAX_WBITS)
+                self._decompressobj = ZLibDecompressor(suppress_deflate_header=True)
             if opcode == WSMsgType.CLOSE:
                 if len(payload) >= 2:
                     close_code = UNPACK_CLOSE_CODE(payload[:2])[0]
@@ -383,8 +384,9 @@ def _feed_data(self, data: bytes) -> Tuple[bool, bytes]:
                     # Decompress process must to be done after all packets
                     # received.
                     if compressed:
+                        assert self._decompressobj is not None
                         self._partial.extend(_WS_DEFLATE_TRAILING)
-                        payload_merged = self._decompressobj.decompress(
+                        payload_merged = self._decompressobj.decompress_sync(
                             self._partial, self._max_msg_size
                         )
                         if self._decompressobj.unconsumed_tail:
@@ -612,16 +614,16 @@ async def _send_frame(
         if (compress or self.compress) and opcode < 8:
             if compress:
                 # Do not set self._compress if compressing is for this frame
-                compressobj = zlib.compressobj(level=zlib.Z_BEST_SPEED, wbits=-compress)
+                compressobj = ZLibCompressor(level=zlib.Z_BEST_SPEED, wbits=-compress)
             else:  # self.compress
                 if not self._compressobj:
-                    self._compressobj = zlib.compressobj(
+                    self._compressobj = ZLibCompressor(
                         level=zlib.Z_BEST_SPEED, wbits=-self.compress
                     )
                 compressobj = self._compressobj
 
-            message = compressobj.compress(message)
-            message = message + compressobj.flush(
+            message = await compressobj.compress(message)
+            message += compressobj.flush(
                 zlib.Z_FULL_FLUSH if self.notakeover else zlib.Z_SYNC_FLUSH
             )
             if message.endswith(_WS_DEFLATE_TRAILING):

aiohttp/http_writer.py

@@ -8,6 +8,7 @@
 
 from .abc import AbstractStreamWriter
 from .base_protocol import BaseProtocol
+from .compression_utils import ZLibCompressor
 from .helpers import NO_EXTENSIONS
 
 __all__ = ("StreamWriter", "HttpVersion", "HttpVersion10", "HttpVersion11")
@@ -43,7 +44,7 @@ def __init__(
         self.output_size = 0
 
         self._eof = False
-        self._compress: Any = None
+        self._compress: Optional[ZLibCompressor] = None
         self._drain_waiter = None
 
         self._on_chunk_sent: _T_OnChunkSent = on_chunk_sent
@@ -63,8 +64,7 @@ def enable_chunking(self) -> None:
     def enable_compression(
         self, encoding: str = "deflate", strategy: int = zlib.Z_DEFAULT_STRATEGY
     ) -> None:
-        zlib_mode = 16 + zlib.MAX_WBITS if encoding == "gzip" else zlib.MAX_WBITS
-        self._compress = zlib.compressobj(wbits=zlib_mode, strategy=strategy)
+        self._compress = ZLibCompressor(encoding=encoding, strategy=strategy)
 
     def _write(self, chunk: bytes) -> None:
         size = len(chunk)
@@ -93,7 +93,7 @@ async def write(
                 chunk = chunk.cast("c")
 
         if self._compress is not None:
-            chunk = self._compress.compress(chunk)
+            chunk = await self._compress.compress(chunk)
             if not chunk:
                 return
 
@@ -138,9 +138,9 @@ async def write_eof(self, chunk: bytes = b"") -> None:
 
         if self._compress:
             if chunk:
-                chunk = self._compress.compress(chunk)
+                chunk = await self._compress.compress(chunk)
 
-            chunk = chunk + self._compress.flush()
+            chunk += self._compress.flush()
             if chunk and self.chunked:
                 chunk_len = ("%x\r\n" % len(chunk)).encode("ascii")
                 chunk = chunk_len + chunk + b"\r\n0\r\n\r\n"