compression.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119
  1. """Compression utilities."""
  2. from __future__ import absolute_import, unicode_literals
  3. from kombu.utils.encoding import ensure_bytes
  4. import bz2
  5. import zlib
  6. try:
  7. import lzma
  8. except ImportError: # pragma: no cover
  9. # TODO: Drop fallback to backports once we drop Python 2.7 support
  10. try:
  11. from backports import lzma
  12. except ImportError: # pragma: no cover
  13. lzma = None
  14. _aliases = {}
  15. _encoders = {}
  16. _decoders = {}
  17. __all__ = ('register', 'encoders', 'get_encoder',
  18. 'get_decoder', 'compress', 'decompress')
  19. def register(encoder, decoder, content_type, aliases=None):
  20. """Register new compression method.
  21. Arguments:
  22. encoder (Callable): Function used to compress text.
  23. decoder (Callable): Function used to decompress previously
  24. compressed text.
  25. content_type (str): The mime type this compression method
  26. identifies as.
  27. aliases (Sequence[str]): A list of names to associate with
  28. this compression method.
  29. """
  30. _encoders[content_type] = encoder
  31. _decoders[content_type] = decoder
  32. if aliases:
  33. _aliases.update((alias, content_type) for alias in aliases)
  34. def encoders():
  35. """Return a list of available compression methods."""
  36. return list(_encoders)
  37. def get_encoder(t):
  38. """Get encoder by alias name."""
  39. t = _aliases.get(t, t)
  40. return _encoders[t], t
  41. def get_decoder(t):
  42. """Get decoder by alias name."""
  43. return _decoders[_aliases.get(t, t)]
  44. def compress(body, content_type):
  45. """Compress text.
  46. Arguments:
  47. body (AnyStr): The text to compress.
  48. content_type (str): mime-type of compression method to use.
  49. """
  50. encoder, content_type = get_encoder(content_type)
  51. return encoder(ensure_bytes(body)), content_type
  52. def decompress(body, content_type):
  53. """Decompress compressed text.
  54. Arguments:
  55. body (AnyStr): Previously compressed text to uncompress.
  56. content_type (str): mime-type of compression method used.
  57. """
  58. return get_decoder(content_type)(body)
  59. register(zlib.compress,
  60. zlib.decompress,
  61. 'application/x-gzip', aliases=['gzip', 'zlib'])
  62. register(bz2.compress,
  63. bz2.decompress,
  64. 'application/x-bz2', aliases=['bzip2', 'bzip'])
  65. try:
  66. import brotli
  67. except ImportError: # pragma: no cover
  68. pass
  69. else:
  70. register(brotli.compress,
  71. brotli.decompress,
  72. 'application/x-brotli', aliases=['brotli'])
  73. # TODO: Drop condition once we drop Python 2.7 support
  74. if lzma: # pragma: no cover
  75. register(lzma.compress,
  76. lzma.decompress,
  77. 'application/x-lzma', aliases=['lzma', 'xz'])
  78. try:
  79. import zstandard as zstd
  80. except ImportError: # pragma: no cover
  81. pass
  82. else:
  83. def zstd_compress(body):
  84. c = zstd.ZstdCompressor()
  85. return c.compress(body)
  86. def zstd_decompress(body):
  87. d = zstd.ZstdDecompressor()
  88. return d.decompress(body)
  89. register(zstd_compress,
  90. zstd_decompress,
  91. 'application/zstd', aliases=['zstd', 'zstandard'])