snappy.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305
  1. #!/usr/bin/env python
  2. #
  3. # Copyright (c) 2011, Andres Moreira <andres@andresmoreira.com>
  4. # 2011, Felipe Cruz <felipecruz@loogica.net>
  5. # 2012, JT Olds <jt@spacemonkey.com>
  6. # All rights reserved.
  7. #
  8. # Redistribution and use in source and binary forms, with or without
  9. # modification, are permitted provided that the following conditions are met:
  10. # * Redistributions of source code must retain the above copyright
  11. # notice, this list of conditions and the following disclaimer.
  12. # * Redistributions in binary form must reproduce the above copyright
  13. # notice, this list of conditions and the following disclaimer in the
  14. # documentation and/or other materials provided with the distribution.
  15. # * Neither the name of the authors nor the
  16. # names of its contributors may be used to endorse or promote products
  17. # derived from this software without specific prior written permission.
  18. #
  19. # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  20. # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  21. # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  22. # ARE DISCLAIMED. IN NO EVENT SHALL ANDRES MOREIRA BE LIABLE FOR ANY DIRECT,
  23. # INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
  24. # (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
  25. # LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
  26. # ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  27. # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  28. # SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  29. #
  30. """python-snappy
  31. Python library for the snappy compression library from Google.
  32. Expected usage like:
  33. import snappy
  34. compressed = snappy.compress("some data")
  35. assert "some data" == snappy.uncompress(compressed)
  36. """
  37. import struct
  38. from _snappy import CompressError, CompressedLengthError, \
  39. InvalidCompressedInputError, UncompressError, \
  40. compress, decompress, isValidCompressed, uncompress, \
  41. _crc32c
  42. _CHUNK_MAX = 65536
  43. _STREAM_TO_STREAM_BLOCK_SIZE = _CHUNK_MAX
  44. _STREAM_IDENTIFIER = b"sNaPpY"
  45. _COMPRESSED_CHUNK = 0x00
  46. _UNCOMPRESSED_CHUNK = 0x01
  47. _IDENTIFIER_CHUNK = 0xff
  48. _RESERVED_UNSKIPPABLE = (0x02, 0x80) # chunk ranges are [inclusive, exclusive)
  49. _RESERVED_SKIPPABLE = (0x80, 0xff)
  50. # the minimum percent of bytes compression must save to be enabled in automatic
  51. # mode
  52. _COMPRESSION_THRESHOLD = .125
  53. def _masked_crc32c(data):
  54. # see the framing format specification
  55. crc = _crc32c(data)
  56. return (((crc >> 15) | (crc << 17)) + 0xa282ead8) & 0xffffffff
  57. _compress = compress
  58. _uncompress = uncompress
  59. class StreamCompressor(object):
  60. """This class implements the compressor-side of the proposed Snappy framing
  61. format, found at
  62. http://code.google.com/p/snappy/source/browse/trunk/framing_format.txt
  63. ?spec=svn68&r=71
  64. This class matches the interface found for the zlib module's compression
  65. objects (see zlib.compressobj), but also provides some additions, such as
  66. the snappy framing format's ability to intersperse uncompressed data.
  67. Keep in mind that this compressor object does no buffering for you to
  68. appropriately size chunks. Every call to StreamCompressor.compress results
  69. in a unique call to the underlying snappy compression method.
  70. """
  71. __slots__ = ["_header_chunk_written"]
  72. def __init__(self):
  73. self._header_chunk_written = False
  74. def add_chunk(self, data, compress=None):
  75. """Add a chunk containing 'data', returning a string that is framed and
  76. (optionally, default) compressed. This data should be concatenated to
  77. the tail end of an existing Snappy stream. In the absence of any
  78. internal buffering, no data is left in any internal buffers, and so
  79. unlike zlib.compress, this method returns everything.
  80. If compress is None, compression is determined automatically based on
  81. snappy's performance. If compress == True, compression always happens,
  82. and if compress == False, compression never happens.
  83. """
  84. if not self._header_chunk_written:
  85. self._header_chunk_written = True
  86. out = [struct.pack("<L", _IDENTIFIER_CHUNK +
  87. (len(_STREAM_IDENTIFIER) << 8)),
  88. _STREAM_IDENTIFIER]
  89. else:
  90. out = []
  91. for i in range(0, len(data), _CHUNK_MAX):
  92. chunk = data[i:i + _CHUNK_MAX]
  93. crc = _masked_crc32c(chunk)
  94. if compress is None:
  95. compressed_chunk = _compress(chunk)
  96. if (len(compressed_chunk) <=
  97. (1 - _COMPRESSION_THRESHOLD) * len(chunk)):
  98. chunk = compressed_chunk
  99. chunk_type = _COMPRESSED_CHUNK
  100. else:
  101. chunk_type = _UNCOMPRESSED_CHUNK
  102. compressed_chunk = None
  103. elif compress:
  104. chunk = _compress(chunk)
  105. chunk_type = _COMPRESSED_CHUNK
  106. else:
  107. chunk_type = _UNCOMPRESSED_CHUNK
  108. out.append(struct.pack("<LL", chunk_type + ((len(chunk) + 4) << 8),
  109. crc))
  110. out.append(chunk)
  111. return b"".join(out)
  112. def compress(self, data):
  113. """This method is simply an alias for compatibility with zlib
  114. compressobj's compress method.
  115. """
  116. return self.add_chunk(data)
  117. def flush(self, mode=None):
  118. """This method does nothing and only exists for compatibility with
  119. the zlib compressobj
  120. """
  121. pass
  122. def copy(self):
  123. """This method exists for compatibility with the zlib compressobj.
  124. """
  125. copy = StreamCompressor()
  126. copy._header_chunk_written = self._header_chunk_written
  127. return copy
  128. class StreamDecompressor(object):
  129. """This class implements the decompressor-side of the proposed Snappy
  130. framing format, found at
  131. http://code.google.com/p/snappy/source/browse/trunk/framing_format.txt
  132. ?spec=svn68&r=71
  133. This class matches a subset of the interface found for the zlib module's
  134. decompression objects (see zlib.decompressobj). Specifically, it currently
  135. implements the decompress method without the max_length option, the flush
  136. method without the length option, and the copy method.
  137. """
  138. __slots__ = ["_buf", "_header_found"]
  139. def __init__(self):
  140. self._buf = b""
  141. self._header_found = False
  142. def decompress(self, data):
  143. """Decompress 'data', returning a string containing the uncompressed
  144. data corresponding to at least part of the data in string. This data
  145. should be concatenated to the output produced by any preceding calls to
  146. the decompress() method. Some of the input data may be preserved in
  147. internal buffers for later processing.
  148. """
  149. self._buf += data
  150. uncompressed = []
  151. while True:
  152. if len(self._buf) < 4:
  153. return b"".join(uncompressed)
  154. chunk_type = struct.unpack("<L", self._buf[:4])[0]
  155. size = (chunk_type >> 8)
  156. chunk_type &= 0xff
  157. if not self._header_found:
  158. if (chunk_type != _IDENTIFIER_CHUNK or
  159. size != len(_STREAM_IDENTIFIER)):
  160. raise UncompressError("stream missing snappy identifier")
  161. self._header_found = True
  162. if (_RESERVED_UNSKIPPABLE[0] <= chunk_type and
  163. chunk_type < _RESERVED_UNSKIPPABLE[1]):
  164. raise UncompressError(
  165. "stream received unskippable but unknown chunk")
  166. if len(self._buf) < 4 + size:
  167. return b"".join(uncompressed)
  168. chunk, self._buf = self._buf[4:4 + size], self._buf[4 + size:]
  169. if chunk_type == _IDENTIFIER_CHUNK:
  170. if chunk != _STREAM_IDENTIFIER:
  171. raise UncompressError(
  172. "stream has invalid snappy identifier")
  173. continue
  174. if (_RESERVED_SKIPPABLE[0] <= chunk_type and
  175. chunk_type < _RESERVED_SKIPPABLE[1]):
  176. continue
  177. assert chunk_type in (_COMPRESSED_CHUNK, _UNCOMPRESSED_CHUNK)
  178. crc, chunk = chunk[:4], chunk[4:]
  179. if chunk_type == _COMPRESSED_CHUNK:
  180. chunk = _uncompress(chunk)
  181. if struct.pack("<L", _masked_crc32c(chunk)) != crc:
  182. raise UncompressError("crc mismatch")
  183. uncompressed.append(chunk)
  184. def flush(self):
  185. """All pending input is processed, and a string containing the
  186. remaining uncompressed output is returned. After calling flush(), the
  187. decompress() method cannot be called again; the only realistic action
  188. is to delete the object.
  189. """
  190. if self._buf != b"":
  191. raise UncompressError("chunk truncated")
  192. return b""
  193. def copy(self):
  194. """Returns a copy of the decompression object. This can be used to save
  195. the state of the decompressor midway through the data stream in order
  196. to speed up random seeks into the stream at a future point.
  197. """
  198. copy = StreamDecompressor()
  199. copy._buf, copy._header_found = self._buf, self._header_found
  200. return copy
  201. def stream_compress(src, dst, blocksize=_STREAM_TO_STREAM_BLOCK_SIZE):
  202. """Takes an incoming file-like object and an outgoing file-like object,
  203. reads data from src, compresses it, and writes it to dst. 'src' should
  204. support the read method, and 'dst' should support the write method.
  205. The default blocksize is good for almost every scenario.
  206. """
  207. compressor = StreamCompressor()
  208. while True:
  209. buf = src.read(blocksize)
  210. if not buf: break
  211. buf = compressor.add_chunk(buf)
  212. if buf: dst.write(buf)
  213. def stream_decompress(src, dst, blocksize=_STREAM_TO_STREAM_BLOCK_SIZE):
  214. """Takes an incoming file-like object and an outgoing file-like object,
  215. reads data from src, decompresses it, and writes it to dst. 'src' should
  216. support the read method, and 'dst' should support the write method.
  217. The default blocksize is good for almost every scenario.
  218. """
  219. decompressor = StreamDecompressor()
  220. while True:
  221. buf = src.read(blocksize)
  222. if not buf: break
  223. buf = decompressor.decompress(buf)
  224. if buf: dst.write(buf)
  225. decompressor.flush() # makes sure the stream ended well
  226. def cmdline_main():
  227. """This method is what is run when invoking snappy via the commandline.
  228. Try python -m snappy --help
  229. """
  230. import sys
  231. if (len(sys.argv) < 2 or len(sys.argv) > 4 or "--help" in sys.argv or
  232. "-h" in sys.argv or sys.argv[1] not in ("-c", "-d")):
  233. print("Usage: python -m snappy <-c/-d> [src [dst]]")
  234. print(" -c compress")
  235. print(" -d decompress")
  236. print("output is stdout if dst is omitted or '-'")
  237. print("input is stdin if src and dst are omitted or src is '-'.")
  238. sys.exit(1)
  239. if len(sys.argv) >= 4 and sys.argv[3] != "-":
  240. dst = open(sys.argv[3], "wb")
  241. elif hasattr(sys.stdout, 'buffer'):
  242. dst = sys.stdout.buffer
  243. else:
  244. dst = sys.stdout
  245. if len(sys.argv) >= 3 and sys.argv[2] != "-":
  246. src = open(sys.argv[2], "rb")
  247. elif hasattr(sys.stdin, "buffer"):
  248. src = sys.stdin.buffer
  249. else:
  250. src = sys.stdin
  251. if sys.argv[1] == "-c":
  252. method = stream_compress
  253. else:
  254. method = stream_decompress
  255. method(src, dst)
  256. if __name__ == "__main__":
  257. cmdline_main()