Allow zero-copy decompression by allowing bytes-like input

blosc/blosc_extension.c

@@ -431,37 +431,47 @@ decompress_helper(void * input, size_t nbytes, void * output)
 
 
 PyDoc_STRVAR(decompress_ptr__doc__,
-"decompress_ptr(string, pointer) -- Decompress string into pointer.\n"
+"decompress_ptr(bytes_like, pointer) -- Decompress bytes-like into pointer.\n"
              );
 
 static PyObject *
 PyBlosc_decompress_ptr(PyObject *self, PyObject *args)
 {
   PyObject * pointer;
-  void * input, * output;
-  size_t cbytes, nbytes;
+  Py_buffer input;
+  void * output;
+  size_t nbytes;
 
   /* require a compressed string and a pointer  */
-  if (!PyArg_ParseTuple(args, "s#O:decompress", &input, &cbytes, &pointer))
+  if (!PyArg_ParseTuple(args, "y*O:decompress_ptr", &input, &pointer)){
+    PyBuffer_Release(&input);
     return NULL;
+  }
 
   /*  convert the int or long Python object to a void * */
   output = PyLong_AsVoidPtr(pointer);
-  if (output == NULL)
+  if (output == NULL){
+    PyBuffer_Release(&input);
     return NULL;
+  }
 
   /*  fetch the uncompressed size into nbytes */
-  if (!get_nbytes(input, cbytes, &nbytes))
+  if (!get_nbytes(input.buf, input.len, &nbytes)){
+    PyBuffer_Release(&input);
     return NULL;
+  }
 
   /* do decompression */
-  if (!decompress_helper(input, nbytes, output))
+  if (!decompress_helper(input.buf, nbytes, output)){
+    PyBuffer_Release(&input);
     return NULL;
+  }
 
   /*  Return nbytes as python integer. This is legitimate, because
    *  decompress_helper above has checked that the number of bytes written
    *  was indeed nbytes.
    *  */
+  PyBuffer_Release(&input);
   return PyLong_FromSize_t(nbytes);
 }
 

blosc/test.py

@@ -91,6 +91,34 @@ def test_decompress_input_types(self):
         self.assertEqual(expected, blosc.decompress(bytearray(compressed)))
         self.assertEqual(expected, blosc.decompress(np.array([compressed])))
 
+    def test_decompress_ptr_input_types(self):
+        import numpy as np
+        # assume the expected answer was compressed from bytes
+        expected = b'0123456789'
+        out = np.zeros(len(expected), dtype=np.byte)
+        compressed = blosc.compress(expected, typesize=1)
+
+        # now for all the things that support the buffer interface
+        out[:] = 0  # reset the output array
+        nout = blosc.decompress_ptr(compressed, out.ctypes.data)
+        self.assertEqual(expected, out.tobytes())
+        self.assertEqual(len(expected), nout)  # check that we didn't write too many bytes
+
+        out[:] = 0
+        nout = blosc.decompress_ptr(memoryview(compressed), out.ctypes.data)
+        self.assertEqual(expected, out.tobytes())
+        self.assertEqual(len(expected), nout)
+
+        out[:] = 0
+        nout = blosc.decompress_ptr(bytearray(compressed), out.ctypes.data)
+        self.assertEqual(expected, out.tobytes())
+        self.assertEqual(len(expected), nout)
+
+        out[:] = 0
+        nout = blosc.decompress_ptr(np.frombuffer(compressed, dtype=np.byte), out.ctypes.data)
+        self.assertEqual(expected, out.tobytes())
+        self.assertEqual(len(expected), nout)
+
     def test_decompress_releasegil(self):
         import numpy as np
         # assume the expected answer was compressed from bytes

blosc/toplevel.py

@@ -373,6 +373,13 @@ def _check_bytesobj(bytesobj):
                         "supported as input")
 
 
+def _check_byteslike(bytes_like):
+    try:
+        memoryview(bytes_like)
+    except:
+        raise TypeError("Input type %s must be a bytes-like object that supports Python Buffer Protocol" % type(bytes_like))
+
+
 def _check_input_length(input_name, input_len):
     if input_len > blosc.MAX_BUFFERSIZE:
         raise ValueError("%s cannot be larger than %d bytes" %
@@ -538,15 +545,17 @@ def compress_ptr(address, items, typesize=8, clevel=9, shuffle=blosc.SHUFFLE,
     return _ext.compress_ptr(address, length, typesize, clevel, shuffle, cname)
 
 
-def decompress(bytesobj, as_bytearray=False):
-    """decompress(bytesobj)
+def decompress(bytes_like, as_bytearray=False):
+    """decompress(bytes_like)
 
-    Decompresses a bytesobj compressed object.
+    Decompresses a bytes-like compressed object.
 
     Parameters
     ----------
-    bytesobj : str / bytes
-        The data to be decompressed.
+    bytes_like : bytes-like object
+        The data to be decompressed.  Must be a bytes-like object
+        that supports the Python Buffer Protocol, like bytes, bytearray,
+        memoryview, or numpy.ndarray.
     as_bytearray : bool, optional
         If this flag is True then the return type will be a bytearray object
         instead of a bytesobject.
@@ -561,7 +570,7 @@ def decompress(bytesobj, as_bytearray=False):
     Raises
     ------
     TypeError
-        If bytesobj is not of type bytes or string.
+        If bytes_like does not support Buffer Protocol
 
     Examples
     --------
@@ -583,18 +592,20 @@ def decompress(bytesobj, as_bytearray=False):
 
     """
 
-    return _ext.decompress(bytesobj, as_bytearray)
+    return _ext.decompress(bytes_like, as_bytearray)
 
 
-def decompress_ptr(bytesobj, address):
-    """decompress_ptr(bytesobj, address)
+def decompress_ptr(bytes_like, address):
+    """decompress_ptr(bytes_like, address)
 
-    Decompresses a bytesobj compressed object into the memory at address.
+    Decompresses a bytes_like compressed object into the memory at address.
 
     Parameters
     ----------
-    bytesobj : str / bytes
-        The data to be decompressed.
+    bytes_like : bytes-like object
+        The data to be decompressed. Must be a bytes-like object
+        that supports the Python Buffer Protocol, like bytes, bytearray,
+        memoryview, or numpy.ndarray.
     address : int or long
         The address at which to write the decompressed data
 
@@ -653,10 +664,10 @@ def decompress_ptr(bytesobj, address):
 
     """
 
-    _check_bytesobj(bytesobj)
+    _check_byteslike(bytes_like)
     _check_address(address)
 
-    return _ext.decompress_ptr(bytesobj, address)
+    return _ext.decompress_ptr(bytes_like, address)
 
 
 def pack_array(array, clevel=9, shuffle=blosc.SHUFFLE, cname='blosclz'):