mirror of
https://github.com/hanzo-ds/connect.git
synced 2026-10-30 20:31:24 +00:00
e29f9fd738
* refactor(driver-utils): add type annotations to driver utility modules * refactor(driver-utils): widen QueryContext.parameters to accept Sequence | dict | None * silence macOS-only mypy error, note query_id change --------- Signed-off-by: Joe Spadola <[email protected]> Co-authored-by: Joe S <[email protected]>
74 lines
1.7 KiB
Python
74 lines
1.7 KiB
Python
import zlib
|
|
|
|
import lz4
|
|
import lz4.frame
|
|
import zstandard
|
|
|
|
try:
|
|
import brotli
|
|
except ImportError:
|
|
brotli = None
|
|
|
|
|
|
available_compression = ["lz4", "zstd"]
|
|
|
|
if brotli:
|
|
available_compression.append("br")
|
|
available_compression.extend(["gzip", "deflate"])
|
|
|
|
comp_map: dict[str, "Compressor | type[Compressor]"] = {}
|
|
|
|
|
|
class Compressor:
|
|
def __init_subclass__(cls, tag: str, thread_safe: bool = True):
|
|
comp_map[tag] = cls() if thread_safe else cls
|
|
|
|
def compress_block(self, block) -> bytes | bytearray:
|
|
return block
|
|
|
|
def flush(self):
|
|
pass
|
|
|
|
|
|
class GzipCompressor(Compressor, tag="gzip", thread_safe=False):
|
|
def __init__(self, level: int = 6, wbits: int = 31):
|
|
self.zlib_obj = zlib.compressobj(level=level, wbits=wbits)
|
|
|
|
def compress_block(self, block):
|
|
return self.zlib_obj.compress(block)
|
|
|
|
def flush(self):
|
|
return self.zlib_obj.flush()
|
|
|
|
|
|
class Lz4Compressor(Compressor, tag="lz4", thread_safe=False):
|
|
def __init__(self):
|
|
self.comp = lz4.frame.LZ4FrameCompressor()
|
|
|
|
def compress_block(self, block):
|
|
output = self.comp.begin(len(block))
|
|
output += self.comp.compress(block)
|
|
return output + self.comp.flush()
|
|
|
|
|
|
class ZstdCompressor(Compressor, tag="zstd"):
|
|
def compress_block(self, block):
|
|
return zstandard.compress(block)
|
|
|
|
|
|
class BrotliCompressor(Compressor, tag="br"):
|
|
def compress_block(self, block):
|
|
return brotli.compress(block)
|
|
|
|
|
|
null_compressor = Compressor()
|
|
|
|
|
|
def get_compressor(compression: str | None) -> Compressor:
|
|
if not compression:
|
|
return null_compressor
|
|
comp = comp_map[compression]
|
|
if isinstance(comp, Compressor):
|
|
return comp
|
|
return comp()
|