@@ -274,7 +274,7 @@ _zstd_train_dict_impl(PyObject *module, PyBytesObject *samples_bytes,
274274 PyObject * samples_sizes , Py_ssize_t dict_size )
275275/*[clinic end generated code: output=8e87fe43935e8f77 input=d20dedb21c72cb62]*/
276276{
277- PyObject * dst_dict_bytes = NULL ;
277+ PyBytesWriter * dst_dict_bytes = NULL ;
278278 size_t * chunk_sizes = NULL ;
279279 Py_ssize_t chunks_number ;
280280 size_t zstd_ret ;
@@ -294,13 +294,13 @@ _zstd_train_dict_impl(PyObject *module, PyBytesObject *samples_bytes,
294294 }
295295
296296 /* Allocate dict buffer */
297- dst_dict_bytes = PyBytes_FromStringAndSize ( NULL , dict_size );
297+ dst_dict_bytes = PyBytesWriter_Create ( dict_size );
298298 if (dst_dict_bytes == NULL ) {
299299 goto error ;
300300 }
301301
302302 /* Train the dictionary */
303- char * dst_dict_buffer = PyBytes_AS_STRING (dst_dict_bytes );
303+ char * dst_dict_buffer = PyBytesWriter_GetData (dst_dict_bytes );
304304 const char * samples_buffer = PyBytes_AS_STRING (samples_bytes );
305305 Py_BEGIN_ALLOW_THREADS
306306 zstd_ret = ZDICT_trainFromBuffer (dst_dict_buffer , dict_size ,
@@ -315,19 +315,15 @@ _zstd_train_dict_impl(PyObject *module, PyBytesObject *samples_bytes,
315315 goto error ;
316316 }
317317
318- /* Resize dict_buffer */
319- if (_PyBytes_Resize (& dst_dict_bytes , zstd_ret ) < 0 ) {
320- goto error ;
321- }
318+ PyMem_Free (chunk_sizes );
322319
323- goto success ;
320+ /* Resize dict_buffer */
321+ return PyBytesWriter_FinishWithSize (dst_dict_bytes , zstd_ret );
324322
325323error :
326- Py_CLEAR (dst_dict_bytes );
327-
328- success :
324+ PyBytesWriter_Discard (dst_dict_bytes );
329325 PyMem_Free (chunk_sizes );
330- return dst_dict_bytes ;
326+ return NULL ;
331327}
332328
333329/*[clinic input]
0 commit comments