@@ -8,8 +8,13 @@ import threading
8
8
9
9
10
10
from numpy cimport ndarray
11
- from cpython.bytes cimport PyBytes_AsString, PyBytes_FromStringAndSize
12
- from cpython.mem cimport PyMem_Malloc, PyMem_Free
11
+ from cpython.ref cimport PyObject
12
+ from cpython.buffer cimport PyObject_GetBuffer, PyBuffer_Release, \
13
+ PyBUF_ANY_CONTIGUOUS
14
+
15
+
16
+ cdef extern from " Python.h" :
17
+ int PyByteArray_Resize(PyObject * bytearray, Py_ssize_t len )
13
18
14
19
15
20
from zarr.compat import PY2, text_type
@@ -70,12 +75,12 @@ def set_nthreads(int nthreads):
70
75
return blosc_set_nthreads(nthreads)
71
76
72
77
73
- def decompress (bytes cdata , ndarray array ):
78
+ def decompress (cdata , ndarray array ):
74
79
""" Decompress data into a numpy array.
75
80
76
81
Parameters
77
82
----------
78
- cdata : bytes
83
+ cdata : bytes-like
79
84
Compressed data, including blosc header.
80
85
array : ndarray
81
86
Numpy array to decompress into.
@@ -90,11 +95,15 @@ def decompress(bytes cdata, ndarray array):
90
95
int ret
91
96
char * source
92
97
char * dest
93
- size_t nbytes
98
+ Py_buffer source_buffer
99
+ Py_buffer dest_buffer
100
+ Py_ssize_t nbytes
94
101
95
102
# setup
96
- source = PyBytes_AsString(cdata)
97
- dest = array.data
103
+ PyObject_GetBuffer(cdata, & source_buffer, PyBUF_ANY_CONTIGUOUS)
104
+ source = < char * > source_buffer.buf
105
+ PyObject_GetBuffer(array, & dest_buffer, PyBUF_ANY_CONTIGUOUS)
106
+ dest = < char * > dest_buffer.buf
98
107
nbytes = array.nbytes
99
108
100
109
# perform decompression
@@ -106,6 +115,10 @@ def decompress(bytes cdata, ndarray array):
106
115
with nogil:
107
116
ret = blosc_decompress_ctx(source, dest, nbytes, 1 )
108
117
118
+ # release buffers
119
+ PyBuffer_Release(& source_buffer)
120
+ PyBuffer_Release(& dest_buffer)
121
+
109
122
# handle errors
110
123
if ret <= 0 :
111
124
raise RuntimeError (' error during blosc decompression: %d ' % ret)
@@ -135,16 +148,21 @@ def compress(ndarray array, char* cname, int clevel, int shuffle):
135
148
cdef:
136
149
char * source
137
150
char * dest
138
- size_t nbytes, cbytes, itemsize
139
- bytes cdata
151
+ bytearray cdata
152
+ Py_buffer source_buffer
153
+ Py_buffer dest_buffer
154
+ Py_ssize_t nbytes, cbytes, itemsize
140
155
141
- # obtain reference to underlying buffer
142
- source = array.data
156
+ # setup source
157
+ PyObject_GetBuffer(array, & source_buffer, PyBUF_ANY_CONTIGUOUS)
158
+ source = < char * > source_buffer.buf
143
159
144
- # allocate memory for compressed data
160
+ # setup destination
145
161
nbytes = array.nbytes
146
162
itemsize = array.dtype.itemsize
147
- dest = < char * > PyMem_Malloc(nbytes + BLOSC_MAX_OVERHEAD)
163
+ cdata = bytearray(nbytes + BLOSC_MAX_OVERHEAD)
164
+ PyObject_GetBuffer(cdata, & dest_buffer, PyBUF_ANY_CONTIGUOUS)
165
+ dest = < char * > dest_buffer.buf
148
166
149
167
# perform compression
150
168
if _get_use_threads():
@@ -163,15 +181,16 @@ def compress(ndarray array, char* cname, int clevel, int shuffle):
163
181
nbytes + BLOSC_MAX_OVERHEAD, cname,
164
182
0 , 1 )
165
183
184
+ # release buffers
185
+ PyBuffer_Release(& source_buffer)
186
+ PyBuffer_Release(& dest_buffer)
187
+
166
188
# check compression was successful
167
189
if cbytes <= 0 :
168
190
raise RuntimeError (' error during blosc compression: %d ' % cbytes)
169
191
170
- # store as bytes
171
- cdata = PyBytes_FromStringAndSize(dest, cbytes)
172
-
173
- # release memory
174
- PyMem_Free(dest)
192
+ # resize after compression
193
+ PyByteArray_Resize(< PyObject * > cdata, cbytes)
175
194
176
195
return cdata
177
196
0 commit comments