@@ -154,8 +154,8 @@ static PyObject *py_zstd_uncompress(PyObject* self, PyObject *args)
154154 const char * source , * src ;
155155 Py_ssize_t source_size , ss , seek_frame ;
156156 uint64_t dest_size , frame_size ;
157- char error = 0 ;
158- size_t cSize ;
157+ char error = 0 , streamed = 0 ;
158+ size_t cSize = 0 , processed = 0 ;
159159
160160#if PY_MAJOR_VERSION >= 3
161161 if (!PyArg_ParseTuple (args , "y#" , & source , & source_size ))
@@ -166,10 +166,15 @@ static PyObject *py_zstd_uncompress(PyObject* self, PyObject *args)
166166#endif
167167
168168 dest_size = (uint64_t ) ZSTD_getFrameContentSize (source , source_size );
169- if (dest_size == ZSTD_CONTENTSIZE_UNKNOWN || dest_size == ZSTD_CONTENTSIZE_ERROR ) {
169+ if (dest_size == ZSTD_CONTENTSIZE_ERROR ) {
170170 PyErr_Format (ZstdError , "Input data invalid or missing content size in frame header." );
171171 return NULL ;
172- }
172+ } else if (dest_size == ZSTD_CONTENTSIZE_UNKNOWN ) {
173+ // probably streamed data
174+ streamed = 1 ;
175+ dest_size = ZSTD_DStreamOutSize ();
176+ } else {
177+ // known block
173178
174179 // Find real dest_size across multiple frames
175180 ss = source_size ;
@@ -185,15 +190,37 @@ static PyObject *py_zstd_uncompress(PyObject* self, PyObject *args)
185190 if (ZSTD_isError (frame_size )) break ;
186191 dest_size += frame_size ;
187192 }
188-
193+ }
189194 result = PyBytes_FromStringAndSize (NULL , dest_size );
190195
191196 if (result != NULL ) {
192197 char * dest = PyBytes_AS_STRING (result );
193198
194199 Py_BEGIN_ALLOW_THREADS
195- // get real dest_size
196- cSize = ZSTD_decompress (dest , dest_size , source , source_size );
200+ if (streamed ) {
201+ ZSTD_DStream * zds ;
202+ zds = ZSTD_createDStream ();
203+ // buffers create and decompress
204+ ZSTD_initDStream (zds );
205+ ZSTD_outBuffer out ;
206+ ZSTD_inBuffer in ;
207+ in .src = source ;
208+ in .pos = 0 ;
209+ in .size = source_size ;
210+ out .dst = dest ;
211+ out .pos = 0 ;
212+ out .size = dest_size ;
213+ processed = ZSTD_decompressStream (zds , & out , & in );
214+ if (processed == 0 ) {
215+ cSize = out .pos ;
216+ if (cSize ) dest_size = cSize ;
217+ }
218+ ZSTD_freeDStream (zds );
219+ }
220+ else {
221+ cSize = ZSTD_decompress (dest , dest_size , source , source_size );
222+ }
223+
197224 Py_END_ALLOW_THREADS
198225
199226 if (ZSTD_isError (cSize )) {
@@ -240,9 +267,9 @@ static PyObject *py_zstd_check(PyObject* self, PyObject *args)
240267{
241268 UNUSED (self );
242269 //PyObject *result;
243- const char * source , * src ;
244- Py_ssize_t source_size , ss , seek_frame ;
245- uint64_t dest_size , frame_size , error = 0 ;
270+ const char * source ;
271+ Py_ssize_t source_size ;
272+ uint64_t dest_size /* , error=0*/ ;
246273 //char error = 0;
247274 //size_t cSize;
248275
@@ -255,29 +282,15 @@ static PyObject *py_zstd_check(PyObject* self, PyObject *args)
255282#endif
256283
257284 dest_size = (uint64_t ) ZSTD_getFrameContentSize (source , source_size );
258- if (dest_size == ZSTD_CONTENTSIZE_UNKNOWN || dest_size == ZSTD_CONTENTSIZE_ERROR ) {
285+ if (dest_size == ZSTD_CONTENTSIZE_ERROR ) {
259286 //PyErr_Format(ZstdError, "Input data invalid or missing content size in frame header.");
260287 return Py_BuildValue ("i" , 0 );
261- } else {
262-
263- // Find real dest_size across multiple frames
264- ss = source_size ;
265- seek_frame = ss - 1 ;
266- src = source ;
267- while (seek_frame < ss ) {
268- seek_frame = ZSTD_findFrameCompressedSize (src , ss );
269- if (ZSTD_isError (seek_frame )) { error = 1 ; break ;}
270- src += seek_frame ;
271- ss -= seek_frame ;
272- if (ss <=0 ) break ;
273- frame_size = (uint64_t ) ZSTD_getFrameContentSize (src , ss );
274- if (ZSTD_isError (frame_size )) { error = 1 ; break ;}
275- dest_size += frame_size ;
276- }
288+ } else if (dest_size == ZSTD_CONTENTSIZE_UNKNOWN ) {
289+ // content valid, just streamed
290+ //dest_size =/* ZSTD_BLOCKSIZE_MAX;*/ ZSTD_DStreamOutSize();
291+ return Py_BuildValue ("i" , 2 );
277292 }
278- if (error ) return Py_BuildValue ("i" , -1 );
279- /*if (ss<=0)
280- return Py_BuildValue("i", 0);*/
293+ //if (error) return Py_BuildValue("i", -1);
281294 return Py_BuildValue ("i" , 1 );
282295}
283296
0 commit comments