cext: accept explicit None value in ZstdCompressor/ZstdDecompressor constructors

Closes #153.

Author: Rodrigo Zhou

c-ext/compressor.c

@@ -97,17 +97,16 @@ static int ZstdCompressor_init(ZstdCompressor *self, PyObject *args,
                              NULL};
 
     int level = 3;
-    ZstdCompressionDict *dict = NULL;
-    ZstdCompressionParametersObject *params = NULL;
+    PyObject *dict = NULL;
+    PyObject *params = NULL;
     PyObject *writeChecksum = NULL;
     PyObject *writeContentSize = NULL;
     PyObject *writeDictID = NULL;
     int threads = 0;
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|iO!O!OOOi:ZstdCompressor",
-                                     kwlist, &level, &ZstdCompressionDictType,
-                                     &dict, &ZstdCompressionParametersType,
-                                     &params, &writeChecksum, &writeContentSize,
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|iOOOOOi:ZstdCompressor",
+                                     kwlist, &level, &dict, &params,
+                                     &writeChecksum, &writeContentSize,
                                      &writeDictID, &threads)) {
         return -1;
     }
@@ -122,6 +121,41 @@ static int ZstdCompressor_init(ZstdCompressor *self, PyObject *args,
         threads = cpu_count();
     }
 
+    if (dict) {
+        if (dict == Py_None) {
+            dict = NULL;
+        }
+        else if (!PyObject_IsInstance(dict,
+                                      (PyObject *)&ZstdCompressionDictType)) {
+            PyErr_Format(PyExc_TypeError,
+                         "dict_data must be zstd.ZstdCompressionDict");
+            return -1;
+        }
+    }
+
+    if (params) {
+        if (params == Py_None) {
+            params = NULL;
+        }
+        else if (!PyObject_IsInstance(
+                     params, (PyObject *)&ZstdCompressionParametersType)) {
+            PyErr_Format(
+                PyExc_TypeError,
+                "compression_params must be zstd.ZstdCompressionParameters");
+            return -1;
+        }
+    }
+
+    if (writeChecksum == Py_None) {
+        writeChecksum = NULL;
+    }
+    if (writeContentSize == Py_None) {
+        writeContentSize = NULL;
+    }
+    if (writeDictID == Py_None) {
+        writeDictID = NULL;
+    }
+
     /* We create a ZSTD_CCtx for reuse among multiple operations to reduce the
        overhead of each compression operation. */
     self->cctx = ZSTD_createCCtx();
@@ -166,7 +200,8 @@ static int ZstdCompressor_init(ZstdCompressor *self, PyObject *args,
     }
 
     if (params) {
-        if (set_parameters(self->params, params)) {
+        if (set_parameters(self->params,
+                           (ZstdCompressionParametersObject *)params)) {
             return -1;
         }
     }
@@ -199,7 +234,7 @@ static int ZstdCompressor_init(ZstdCompressor *self, PyObject *args,
     }
 
     if (dict) {
-        self->dict = dict;
+        self->dict = (ZstdCompressionDict *)dict;
         Py_INCREF(dict);
     }
 

c-ext/decompressor.c

@@ -29,7 +29,8 @@ int ensure_dctx(ZstdDecompressor *decompressor, int loadDict) {
         }
     }
 
-    zresult = ZSTD_DCtx_setParameter(decompressor->dctx, ZSTD_d_format, decompressor->format);
+    zresult = ZSTD_DCtx_setParameter(decompressor->dctx, ZSTD_d_format,
+                                     decompressor->format);
     if (ZSTD_isError(zresult)) {
         PyErr_Format(ZstdError, "unable to set decoding format: %s",
                      ZSTD_getErrorName(zresult));
@@ -58,19 +59,30 @@ static int Decompressor_init(ZstdDecompressor *self, PyObject *args,
                              PyObject *kwargs) {
     static char *kwlist[] = {"dict_data", "max_window_size", "format", NULL};
 
-    ZstdCompressionDict *dict = NULL;
+    PyObject *dict = NULL;
     Py_ssize_t maxWindowSize = 0;
     ZSTD_format_e format = ZSTD_f_zstd1;
 
     self->dctx = NULL;
     self->dict = NULL;
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O!nI:ZstdDecompressor",
-                                     kwlist, &ZstdCompressionDictType, &dict,
-                                     &maxWindowSize, &format)) {
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|OnI:ZstdDecompressor",
+                                     kwlist, &dict, &maxWindowSize, &format)) {
         return -1;
     }
 
+    if (dict) {
+        if (dict == Py_None) {
+            dict = NULL;
+        }
+        else if (!PyObject_IsInstance(dict,
+                                      (PyObject *)&ZstdCompressionDictType)) {
+            PyErr_Format(PyExc_TypeError,
+                         "dict_data must be zstd.ZstdCompressionDict");
+            return -1;
+        }
+    }
+
     self->dctx = ZSTD_createDCtx();
     if (!self->dctx) {
         PyErr_NoMemory();
@@ -81,7 +93,7 @@ static int Decompressor_init(ZstdDecompressor *self, PyObject *args,
     self->format = format;
 
     if (dict) {
-        self->dict = dict;
+        self->dict = (ZstdCompressionDict *)dict;
         Py_INCREF(dict);
     }
 

docs/news.rst

@@ -95,6 +95,11 @@ Changes
 * ``manylinux2014_aarch64`` wheels are now being produced for CPython 3.6+.
   (#145).
 * Wheels are now being produced for CPython 3.10.
+* Arguments to ``ZstdCompressor()`` and ``ZstdDecompressor()`` are now all
+  optional in the C backend and an explicit ``None`` value is accepted. Before,
+  the C backend wouldn't accept an explicit ``None`` value (but the CFFI
+  backend would). The new behavior should be consistent between the backends.
+  (#153)
 
 0.15.2 (released 2021-02-27)
 ============================

tests/test_compressor_compress.py

@@ -213,3 +213,74 @@ def test_multithreaded_compression_params(self):
         self.assertEqual(
             result, b"\x28\xb5\x2f\xfd\x20\x03\x19\x00\x00\x66\x6f\x6f"
         )
+
+    def test_explicit_default_params(self):
+        cctx = zstd.ZstdCompressor(
+            level=3,
+            dict_data=None,
+            compression_params=None,
+            write_checksum=None,
+            write_content_size=None,
+            write_dict_id=None,
+            threads=0,
+        )
+        result = cctx.compress(b"")
+        self.assertEqual(result, b"\x28\xb5\x2f\xfd\x20\x00\x01\x00\x00")
+
+    def test_compression_params_with_other_params(self):
+        params = zstd.ZstdCompressionParameters.from_level(3)
+        cctx = zstd.ZstdCompressor(
+            level=3,
+            dict_data=None,
+            compression_params=params,
+            write_checksum=None,
+            write_content_size=None,
+            write_dict_id=None,
+            threads=0,
+        )
+        result = cctx.compress(b"")
+        self.assertEqual(result, b"\x28\xb5\x2f\xfd\x20\x00\x01\x00\x00")
+
+        with self.assertRaises(ValueError):
+            cctx = zstd.ZstdCompressor(
+                level=3,
+                dict_data=None,
+                compression_params=params,
+                write_checksum=False,
+                write_content_size=None,
+                write_dict_id=None,
+                threads=0,
+            )
+
+        with self.assertRaises(ValueError):
+            cctx = zstd.ZstdCompressor(
+                level=3,
+                dict_data=None,
+                compression_params=params,
+                write_checksum=None,
+                write_content_size=True,
+                write_dict_id=None,
+                threads=0,
+            )
+
+        with self.assertRaises(ValueError):
+            cctx = zstd.ZstdCompressor(
+                level=3,
+                dict_data=None,
+                compression_params=params,
+                write_checksum=None,
+                write_content_size=None,
+                write_dict_id=True,
+                threads=0,
+            )
+
+        with self.assertRaises(ValueError):
+            cctx = zstd.ZstdCompressor(
+                level=3,
+                dict_data=None,
+                compression_params=params,
+                write_checksum=None,
+                write_content_size=None,
+                write_dict_id=True,
+                threads=2,
+            )

tests/test_decompressor_decompress.py

@@ -164,3 +164,14 @@ def test_max_window_size(self):
             "decompression error: Frame requires too much memory",
         ):
             dctx.decompress(frame, max_output_size=len(source))
+
+    def test_explicit_default_params(self):
+        cctx = zstd.ZstdCompressor(level=1)
+        compressed = cctx.compress(b"foo")
+
+        dctx = zstd.ZstdDecompressor(
+            dict_data=None,
+            max_window_size=0,
+            format=zstd.FORMAT_ZSTD1,
+        )
+        self.assertEqual(dctx.decompress(compressed), b"foo")

zstandard/__init__.pyi

@@ -265,11 +265,11 @@ class ZstdCompressor(object):
     def __init__(
         self,
         level: int = ...,
-        dict_data: ZstdCompressionDict = ...,
-        compression_params: ZstdCompressionParameters = ...,
-        write_checksum: bool = ...,
-        write_content_size: bool = ...,
-        write_dict_id: bool = ...,
+        dict_data: Optional[ZstdCompressionDict] = ...,
+        compression_params: Optional[ZstdCompressionParameters] = ...,
+        write_checksum: Optional[bool] = ...,
+        write_content_size: Optional[bool] = ...,
+        write_dict_id: Optional[bool] = ...,
         threads: int = ...,
     ): ...
     def memory_size(self) -> int: ...
@@ -376,7 +376,7 @@ class ZstdDecompressionWriter(BinaryIO):
 class ZstdDecompressor(object):
     def __init__(
         self,
-        dict_data: ZstdCompressionDict = ...,
+        dict_data: Optional[ZstdCompressionDict] = ...,
         max_window_size: int = ...,
         format: int = ...,
     ): ...