-
-
Notifications
You must be signed in to change notification settings - Fork 2k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
refactor: Extract zlib-related logic into a single module (#7223)
<!-- Thank you for your contribution! --> ## What do these changes do? Addresses issue #7192 Refactors the logic to have the zlib-related stuff concentrated into a single module ## Are there changes in behavior for the user? No ## Related issue number #7192 ## Checklist - [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 <Name> <Surname>. * 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]>
- Loading branch information
3 people
authored
Mar 7, 2023
1 parent
3058c72
commit 3ff81dc
Showing
8 changed files
with
208 additions
and
87 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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"" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.