Merge pull request #77 from khvn26/master

Fix #76: support new buffer API

Author: martindurant

snappy/snappy.py

@@ -127,13 +127,12 @@ def add_chunk(self, data, compress=None):
         snappy's performance. If compress == True, compression always happens,
         and if compress == False, compression never happens.
         """
+        out = bytearray()
         if not self._header_chunk_written:
             self._header_chunk_written = True
-            out = [struct.pack("<L", _IDENTIFIER_CHUNK +
-                                      (len(_STREAM_IDENTIFIER) << 8)),
-                   _STREAM_IDENTIFIER]
-        else:
-            out = []
+            out.extend(struct.pack("<L", _IDENTIFIER_CHUNK +
+                                      (len(_STREAM_IDENTIFIER) << 8)))
+            out.extend(_STREAM_IDENTIFIER)
         for i in range(0, len(data), _CHUNK_MAX):
             chunk = data[i:i + _CHUNK_MAX]
             crc = _masked_crc32c(chunk)
@@ -151,10 +150,10 @@ def add_chunk(self, data, compress=None):
                 chunk_type = _COMPRESSED_CHUNK
             else:
                 chunk_type = _UNCOMPRESSED_CHUNK
-            out.append(struct.pack("<LL", chunk_type + ((len(chunk) + 4) << 8),
+            out.extend(struct.pack("<LL", chunk_type + ((len(chunk) + 4) << 8),
                                    crc))
-            out.append(chunk)
-        return b"".join(out)
+            out.extend(chunk)
+        return bytes(out)
 
     def compress(self, data):
         """This method is simply an alias for compatibility with zlib
@@ -193,7 +192,7 @@ class StreamDecompressor(object):
     __slots__ = ["_buf", "_header_found"]
 
     def __init__(self):
-        self._buf = b""
+        self._buf = bytearray()
         self._header_found = False
 
     @staticmethod
@@ -222,11 +221,11 @@ def decompress(self, data):
         the decompress() method. Some of the input data may be preserved in
         internal buffers for later processing.
         """
-        self._buf += data
-        uncompressed = []
+        self._buf.extend(data)
+        uncompressed = bytearray()
         while True:
             if len(self._buf) < 4:
-                return b"".join(uncompressed)
+                return bytes(uncompressed)
             chunk_type = struct.unpack("<L", self._buf[:4])[0]
             size = (chunk_type >> 8)
             chunk_type &= 0xff
@@ -240,7 +239,7 @@ def decompress(self, data):
                 raise UncompressError(
                     "stream received unskippable but unknown chunk")
             if len(self._buf) < 4 + size:
-                return b"".join(uncompressed)
+                return bytes(uncompressed)
             chunk, self._buf = self._buf[4:4 + size], self._buf[4 + size:]
             if chunk_type == _IDENTIFIER_CHUNK:
                 if chunk != _STREAM_IDENTIFIER:
@@ -256,7 +255,7 @@ def decompress(self, data):
                 chunk = _uncompress(chunk)
             if struct.pack("<L", _masked_crc32c(chunk)) != crc:
                 raise UncompressError("crc mismatch")
-            uncompressed.append(chunk)
+            uncompressed += chunk
 
     def flush(self):
         """All pending input is processed, and a string containing the
@@ -274,7 +273,7 @@ def copy(self):
         to speed up random seeks into the stream at a future point.
         """
         copy = StreamDecompressor()
-        copy._buf, copy._header_found = self._buf, self._header_found
+        copy._buf, copy._header_found = bytearray(self._buf), self._header_found
         return copy
 
 

snappy/snappy_cffi.py

@@ -20,7 +20,7 @@ def prepare(data):
     _out_data = None
     _out_size = None
 
-    _out_data = ffi.new('char[]', data)
+    _out_data = ffi.from_buffer(data)
     _out_size = ffi.cast('size_t', len(data))
 
     return (_out_data, _out_size)
@@ -89,6 +89,6 @@ def isValidCompressed(data):
 decompress = uncompress
 
 def _crc32c(data):
-    c_data = ffi.new('char[]', data)
+    c_data = ffi.from_buffer(data)
     size = ffi.cast('int', len(data))
     return int(lib._crc32c(c_data, size))

snappy/snappymodule.cc

@@ -203,18 +203,22 @@ snappy__is_valid_compressed_buffer(PyObject *self, PyObject *args)
 static PyObject *
 snappy__crc32c(PyObject *self, PyObject *args)
 {
-    const unsigned char * input;
-    int input_size;
+    Py_buffer input;
+    PyObject * result;
 
 #if PY_MAJOR_VERSION >= 3
-    if (!PyArg_ParseTuple(args, "y#", &input, &input_size))
+    if (!PyArg_ParseTuple(args, "y*", &input))
 #else
-    if (!PyArg_ParseTuple(args, "s#", &input, &input_size))
+    if (!PyArg_ParseTuple(args, "s*", &input))
 #endif
         return NULL;
 
-    return PyLong_FromUnsignedLong(
-            crc_finalize(crc_update(crc_init(), input, input_size)));
+    result = PyLong_FromUnsignedLong(
+            crc_finalize(crc_update(crc_init(), (const unsigned char *) input.buf, input.len)));
+
+    PyBuffer_Release(&input);
+
+    return result;
 }
 
 static PyMethodDef snappy_methods[] = {

test_snappy.py

@@ -163,6 +163,18 @@ def test_compression(self):
                 b"\x01\x04\x00\x01" + crc1 + data[:snappy.snappy._CHUNK_MAX] +
                 b"\x01\xff\xff\x00" + crc2 + data[snappy.snappy._CHUNK_MAX:])
 
+        # test that we can compress a memoryview
+        data = memoryview(b"\x01" * (snappy.snappy._CHUNK_MAX * 2 - 5))
+        crc1 = struct.pack("<L",
+                snappy.snappy._masked_crc32c(data[:snappy.snappy._CHUNK_MAX]))
+        self.assertEqual(crc1, b"h#6\x8e")
+        crc2 = struct.pack("<L",
+                snappy.snappy._masked_crc32c(data[snappy.snappy._CHUNK_MAX:]))
+        self.assertEqual(crc2, b"q\x8foE")
+        self.assertEqual(compressor.add_chunk(data, compress=False),
+                b"\x01\x04\x00\x01" + crc1 + data[:snappy.snappy._CHUNK_MAX].tobytes() +
+                b"\x01\xff\xff\x00" + crc2 + data[snappy.snappy._CHUNK_MAX:].tobytes())
+
     def test_decompression(self):
         # test that we check for the initial stream identifier
         data = b"\x01" * 50
@@ -218,6 +230,15 @@ def test_decompression(self):
                      decompressor.decompress(compressed_data[split2:])),
                     uncompressed_data)
 
+        # test that we can decompress a memoryview
+        decompressor = snappy.StreamDecompressor()
+        decompressor.decompress(memoryview(b"\xff\x06\x00\x00sNaPpY"))
+        self.assertEqual(
+                decompressor.copy().decompress(
+                    memoryview(b"\x01\x36\x00\x00" +
+                    struct.pack("<L", snappy.snappy._masked_crc32c(data)) + data)),
+                data)
+
     def test_concatenation(self):
         data1 = os.urandom(snappy.snappy._CHUNK_MAX * 2)
         data2 = os.urandom(4096)