1 files changed, 40 insertions, 13 deletions

diff --git a/Modules/_zstd/compressor.c b/Modules/_zstd/compressor.c
index b92083cf78f..e70eb637b29 100644
--- a/Modules/_zstd/compressor.c
+++ b/Modules/_zstd/compressor.c
@@ -7,24 +7,50 @@ Python module.
 
 /*[clinic input]
 module _zstd
-class _zstd.ZstdCompressor "ZstdCompressor *" "clinic_state()->ZstdCompressor_type"
+class _zstd.ZstdCompressor "ZstdCompressor *" "&zstd_compressor_type_spec"
 [clinic start generated code]*/
-/*[clinic end generated code: output=da39a3ee5e6b4b0d input=875bf614798f80cb]*/
-
+/*[clinic end generated code: output=da39a3ee5e6b4b0d input=7166021db1ef7df8]*/
 
 #ifndef Py_BUILD_CORE_BUILTIN
 #  define Py_BUILD_CORE_MODULE 1
 #endif
 
-#include "_zstdmodule.h"
+#include "Python.h"
 
+#include "_zstdmodule.h"
 #include "buffer.h"
+#include "zstddict.h"
 
+#include <stdbool.h>              // bool
 #include <stddef.h>               // offsetof()
+#include <zstd.h>                 // ZSTD_*()
+
+typedef struct {
+    PyObject_HEAD
+
+    /* Compression context */
+    ZSTD_CCtx *cctx;
+
+    /* ZstdDict object in use */
+    PyObject *dict;
+
+    /* Last mode, initialized to ZSTD_e_end */
+    int last_mode;
+
+    /* (nbWorker >= 1) ? 1 : 0 */
+    int use_multithread;
 
+    /* Compression level */
+    int compression_level;
+
+    /* __init__ has been called, 0 or 1. */
+    bool initialized;
+} ZstdCompressor;
 
 #define ZstdCompressor_CAST(op) ((ZstdCompressor *)op)
 
+#include "clinic/compressor.c.h"
+
 static int
 _zstd_set_c_parameters(ZstdCompressor *self, PyObject *level_or_options,
                        const char *arg_name, const char* arg_type)
@@ -292,10 +318,6 @@ load:
     return 0;
 }
 
-#define clinic_state() (get_zstd_state_from_type(type))
-#include "clinic/compressor.c.h"
-#undef clinic_state
-
 static PyObject *
 _zstd_ZstdCompressor_new(PyTypeObject *type, PyObject *Py_UNUSED(args), PyObject *Py_UNUSED(kwargs))
 {
@@ -305,7 +327,7 @@ _zstd_ZstdCompressor_new(PyTypeObject *type, PyObject *Py_UNUSED(args), PyObject
         goto error;
     }
 
-    self->inited = 0;
+    self->initialized = 0;
     self->dict = NULL;
     self->use_multithread = 0;
 
@@ -372,12 +394,11 @@ _zstd_ZstdCompressor___init___impl(ZstdCompressor *self, PyObject *level,
                                    PyObject *options, PyObject *zstd_dict)
 /*[clinic end generated code: output=215e6c4342732f96 input=9f79b0d8d34c8ef0]*/
 {
-    /* Only called once */
-    if (self->inited) {
-        PyErr_SetString(PyExc_RuntimeError, init_twice_msg);
+    if (self->initialized) {
+        PyErr_SetString(PyExc_RuntimeError, "reinitialization not supported");
         return -1;
     }
-    self->inited = 1;
+    self->initialized = 1;
 
     if (level != Py_None && options != Py_None) {
         PyErr_SetString(PyExc_RuntimeError, "Only one of level or options should be used.");
@@ -488,6 +509,12 @@ error:
     return NULL;
 }
 
+static inline int
+mt_continue_should_break(ZSTD_inBuffer *in, ZSTD_outBuffer *out)
+{
+    return in->size == in->pos && out->size != out->pos;
+}
+
 static PyObject *
 compress_mt_continue_impl(ZstdCompressor *self, Py_buffer *data)
 {