| 1 | 1 | import base64 |
|
| 2 | - | import bz2 |
|
| 3 | - | import gzip |
|
| 4 | 2 | import io |
|
| 5 | 3 | import json |
|
| 6 | - | import lzma |
|
| 4 | + | import importlib |
|
| 5 | + | ||
| 7 | 6 | from typing import TYPE_CHECKING, Any, Callable, Dict, Tuple |
|
| 8 | - | import zlib |
|
| 9 | 7 | ||
| 10 | 8 | import cloudpickle |
|
| 11 | 9 | import pendulum |
| 22 | 20 | "CompressedSerializer", |
|
| 23 | 21 | ) |
|
| 24 | 22 | ||
| 25 | - | COMPRESSION_FORMATS: Dict[str, Tuple[Callable[..., bytes], Callable[..., bytes]]] = { |
|
| 26 | - | "bz2": (bz2.compress, bz2.decompress), |
|
| 27 | - | "gzip": (gzip.compress, gzip.decompress), |
|
| 28 | - | "lzma": (lzma.compress, lzma.decompress), |
|
| 29 | - | "zlib": (zlib.compress, zlib.decompress), |
|
| 30 | - | } |
|
| 31 | - | ||
| 32 | 23 | ||
| 33 | 24 | class Serializer: |
|
| 34 | 25 | """ |
| 270 | 261 | ||
| 271 | 262 | Args: |
|
| 272 | 263 | - serializer (Serializer): the serializer that this serializer wraps |
|
| 273 | - | - format (str): name of the selected pre-defined compression format (bz2, gzip, |
|
| 274 | - | lzma, or zlib) |
|
| 264 | + | - format (str): name of the compression format library. Typically one of the |
|
| 265 | + | python standard compression libraries: bz2, gzip, lzma, or zlib. Attempts |
|
| 266 | + | to import the given format's module and retrieves the compress/decompress |
|
| 267 | + | functions. |
|
| 275 | 268 | - compress (Callable[..., bytes]): the custom compression function |
|
| 276 | 269 | - decompress (Callable[..., bytes]): the custom decompression function |
|
| 277 | 270 | - compress_kwargs (Dict[str, Any]): keyword arguments to be passed to the |
| 297 | 290 | "but not both." |
|
| 298 | 291 | ) |
|
| 299 | 292 | elif format: |
|
| 300 | - | try: |
|
| 301 | - | self._compress, self._decompress = COMPRESSION_FORMATS[format] |
|
| 302 | - | except KeyError as e: |
|
| 303 | - | raise ValueError( |
|
| 304 | - | "`format` must be one of: {}".format(", ".join(COMPRESSION_FORMATS)) |
|
| 305 | - | ) from e |
|
| 293 | + | self._compress, self._decompress = self.compression_from_lib(format) |
|
| 306 | 294 | elif compress and decompress: |
|
| 307 | 295 | self._compress = compress |
|
| 308 | 296 | self._decompress = decompress |
| 351 | 339 | return self._serializer.deserialize( |
|
| 352 | 340 | self._decompress(value, **self._decompress_kwargs) |
|
| 353 | 341 | ) |
|
| 342 | + | ||
| 343 | + | @staticmethod |
|
| 344 | + | def compression_from_lib( |
|
| 345 | + | compression_format: str, |
|
| 346 | + | ) -> Tuple[Callable[..., bytes], Callable[..., bytes]]: |
|
| 347 | + | """ |
|
| 348 | + | Attempt to pull a compression format from a library. Typically one of |
|
| 349 | + | "lzma", "gzip", "zlib", "bz2" |
|
| 350 | + | ||
| 351 | + | Args: |
|
| 352 | + | - compression_format: The compression format/library to load |
|
| 353 | + | ||
| 354 | + | Returns: |
|
| 355 | + | A tuple of functions for compression and decompression |
|
| 356 | + | """ |
|
| 357 | + | common_libs = {"lzma", "gzip", "zlib", "bz2"} |
|
| 358 | + | ||
| 359 | + | # Don't suggest them a format they've just requested |
|
| 360 | + | common_libs.discard(compression_format) |
|
| 361 | + | ||
| 362 | + | try: |
|
| 363 | + | module = importlib.import_module(compression_format) |
|
| 364 | + | except ImportError as exc: |
|
| 365 | + | raise ImportError( |
|
| 366 | + | f"Compression module for {compression_format!r} is not installed. " |
|
| 367 | + | f"Did you mean to use one of {common_libs}?" |
|
| 368 | + | ) from exc |
|
| 369 | + | ||
| 370 | + | try: |
|
| 371 | + | funcs = (module.compress, module.decompress) # type: ignore |
|
| 372 | + | except AttributeError as exc: |
|
| 373 | + | raise ValueError( |
|
| 374 | + | f"Given compression format {compression_format!r} module does not have " |
|
| 375 | + | f"'compress' and 'decompress' attributes. Pass these functions " |
|
| 376 | + | f"manually instead if you intend to use a non-standard library. " |
|
| 377 | + | f"Otherwise, use one of the common compression libraries: {common_libs}" |
|
| 378 | + | ) from exc |
|
| 379 | + | ||
| 380 | + | return funcs |
| Files | Coverage |
|---|---|
| src/prefect | 94.51% |
| Project Totals (108 files) | 94.51% |
| 1 |
comment: |
| 2 |
layout: "header" |
| 3 |
behavior: default |
| 4 |
require_changes: false # if true: only post the comment if coverage changes |
| 5 |
require_base: no # [yes :: must have a base report to post] |
| 6 |
require_head: yes # [yes :: must have a head report to post] |
| 7 |
branches: null # branch names that can post comment |
| 8 |
|
| 9 |
coverage: |
| 10 |
status: |
| 11 |
patch: no |
| 12 |
project: |
| 13 |
default: |
| 14 |
target: 92% |
| 15 |
threshold: 1% |
| 16 |
|
| 17 |
ignore: |
| 18 |
- "src/prefect/tasks/**/*" |
| 19 |
- "src/prefect/environments/**/*" |
madkinsz
