Add as_bytearray=False flag to decompress.

blosc/blosc_extension.c

@@ -365,8 +365,9 @@ PyBlosc_decompress_ptr(PyObject *self, PyObject *args)
 }
 
 PyDoc_STRVAR(decompress__doc__,
-"decompress(string) -- Return decompressed string.\n"
-             );
+"decompress(string, as_bytearray) -- Return decompressed string.\n\n"
+"If as_bytearray is True then the returned data will be a mutable\n"
+"bytearray object instead of bytes");
 
 static PyObject *
 PyBlosc_decompress(PyObject *self, PyObject *args)
@@ -376,21 +377,28 @@ PyBlosc_decompress(PyObject *self, PyObject *args)
     void *input, *output;
     size_t nbytes, cbytes;
     char *format;
+    int as_bytearray;
     /* Accept some kind of input */
     #if PY_MAJOR_VERSION <= 2
+        PyObject *as_bytearray_obj = NULL;
         /* s* : bytes like object including unicode and anything that supports
-         * the buffer interface */
-        format = "s*:decompress";
+         * the buffer interface. We cannot use p in python 2 so we will
+         * create an object to hold the predicate. */
+        if (!PyArg_ParseTuple(args, "s*O:decompress", &view, &as_bytearray_obj))
+            return NULL;
+
+        if ((as_bytearray = PyObject_IsTrue(as_bytearray_obj)) < 0) {
+            /* failed to convert predicate to bool */
+            return NULL;
+        }
     #elif PY_MAJOR_VERSION >= 3
         /* y* :bytes like object EXCLUDING unicode and anything that supports
          * the buffer interface. This is the recommended way to accept binary
          * data in Python 3. */
-        format = "y*:decompress";
+        if (!PyArg_ParseTuple(args, "y*p:decompress", &view, &as_bytearray))
+            return NULL;
     #endif
 
-    if (!PyArg_ParseTuple(args, format, &view))
-      return NULL;
-
     cbytes = view.len;
     input = view.buf;
     /*  fetch the uncompressed size into nbytes */
@@ -399,13 +407,24 @@ PyBlosc_decompress(PyObject *self, PyObject *args)
       return NULL;
     }
 
-    /* Book memory for the result */
-    if (!(result_str = PyBytes_FromStringAndSize(NULL, (Py_ssize_t)nbytes))){
-      PyBuffer_Release(&view);
-      return NULL;
+    #define branch(from_string_and_size, as_string)                     \
+        /* Book memory for the result */                                \
+        if (!(result_str = from_string_and_size(NULL, (Py_ssize_t)nbytes))){ \
+          PyBuffer_Release(&view);                                      \
+          return NULL;                                                  \
+        }                                                               \
+                                                                        \
+        output = as_string(result_str);                                 \
+        (void)NULL
+
+    if (as_bytearray) {
+        branch(PyByteArray_FromStringAndSize, PyByteArray_AS_STRING);
+    }
+    else {
+        branch(PyBytes_FromStringAndSize, PyBytes_AS_STRING);
     }
 
-    output = PyBytes_AS_STRING(result_str);
+    #undef branch
 
     /*  do decompression */
     if (!decompress_helper(input, nbytes, output)){

blosc/test.py

@@ -101,6 +101,28 @@ def test_decompress_input_types(self):
         self.assertEqual(expected, blosc.decompress(bytearray(compressed)))
         self.assertEqual(expected, blosc.decompress(np.array([compressed])))
 
+    def test_decompress_input_types_as_bytearray(self):
+        import numpy as np
+        # assume the expected answer was compressed from bytes
+        expected = bytearray(b'0123456789')
+        compressed = blosc.compress(expected, typesize=1)
+
+        # now for all the things that support the buffer interface
+        if not PY3X:
+            # Python 3 no longer has the buffer
+            self.assertEqual(expected, blosc.decompress(buffer(compressed),
+                                                        as_bytearray=True))
+        if not PY26:
+            # memoryview doesn't exist on Python 2.6
+            self.assertEqual(expected,
+                             blosc.decompress(memoryview(compressed),
+                                              as_bytearray=True))
+
+        self.assertEqual(expected, blosc.decompress(bytearray(compressed),
+                                                    as_bytearray=True))
+        self.assertEqual(expected, blosc.decompress(np.array([compressed]),
+                                                    as_bytearray=True))
+
     def test_compress_exceptions(self):
         s = b'0123456789'
 

blosc/toplevel.py

@@ -421,7 +421,7 @@ def compress_ptr(address, items, typesize, clevel=9, shuffle=True,
     return _ext.compress_ptr(address, length, typesize, clevel, shuffle, cname)
 
 
-def decompress(bytesobj):
+def decompress(bytesobj, as_bytearray=False):
     """decompress(bytesobj)
 
     Decompresses a bytesobj compressed object.
@@ -430,11 +430,16 @@ def decompress(bytesobj):
     ----------
     bytesobj : str / bytes
         The data to be decompressed.
+    as_bytearray : bool, optional
+        If this flag is True then the return type will be a bytearray object
+        instead of a bytesobject.
 
     Returns
     -------
-    out : str / bytes
+    out : str / bytes or bytearray
         The decompressed data in form of a Python str / bytes object.
+        If as_bytearray is True then this will be a bytearray object, otherwise
+        this will be a str/ bytes object.
 
     Raises
     ------
@@ -455,10 +460,13 @@ def decompress(bytesobj):
     True
     >>> b"1"*7 == blosc.decompress(blosc.compress(b"1"*7, 8))
     True
+    >>> type(blosc.decompress(blosc.compress(b"1"*7, 8),
+    ...                                      as_bytearray=True)) is bytearray
+    True
 
     """
 
-    return _ext.decompress(bytesobj)
+    return _ext.decompress(bytesobj, as_bytearray)
 
 
 def decompress_ptr(bytesobj, address):
@@ -637,7 +645,7 @@ def unpack_array(packed_array):
     _check_bytesobj(packed_array)
 
     # First decompress the pickle
-    pickled_array = _ext.decompress(packed_array)
+    pickled_array = _ext.decompress(packed_array, False)
     # ... and unpickle
     array = pickle.loads(pickled_array)