compressionwriter: add closefd argument to control closing underlying…

… stream

The existing behavior of automatically calling close() on the underlying
stream is not desired in all cases. See #76 for discussion.

This commit adds an argument to ZstdCompressionWriter to control whether
the inner stream is closed when the outer stream is closed. This will give
consumers control over how to behave.

We plan to implement this feature for all file-like types. After we're
done, we'll revisit the default behavior.

NEWS.rst

@@ -145,6 +145,9 @@ Changes
 * All built/installed Python modules are now in the ``zstandard``
   package. Previously, there were modules in other packages. (#115)
 * C source code is now automatically formatted with ``clang-format``.
+* ``ZstdCompressor.stream_writer()`` now accepts a ``closefd`` argument
+  to control whether the underlying stream should be closed when the
+  ``ZstdCompressionWriter`` is closed.
 
 0.14.1 (released 2020-12-05)
 ============================

c-ext/compressionwriter.c

@@ -216,7 +216,7 @@ static PyObject *ZstdCompressionWriter_close(ZstdCompressionWriter *self) {
     }
 
     /* Call close on underlying stream as well. */
-    if (PyObject_HasAttrString(self->writer, "close")) {
+    if (self->closefd && PyObject_HasAttrString(self->writer, "close")) {
         return PyObject_CallMethod(self->writer, "close", NULL);
     }
 

c-ext/compressor.c

@@ -799,19 +799,20 @@ PyDoc_STRVAR(
 static ZstdCompressionWriter *ZstdCompressor_stream_writer(ZstdCompressor *self,
                                                            PyObject *args,
                                                            PyObject *kwargs) {
-    static char *kwlist[] = {"writer", "size", "write_size",
-                             "write_return_read", NULL};
+    static char *kwlist[] = {
+        "writer", "size", "write_size", "write_return_read", "closefd", NULL};
 
     PyObject *writer;
     ZstdCompressionWriter *result;
     size_t zresult;
     unsigned long long sourceSize = ZSTD_CONTENTSIZE_UNKNOWN;
     size_t outSize = ZSTD_CStreamOutSize();
     PyObject *writeReturnRead = NULL;
+    PyObject *closefd = NULL;
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|KkO:stream_writer",
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|KkOO:stream_writer",
                                      kwlist, &writer, &sourceSize, &outSize,
-                                     &writeReturnRead)) {
+                                     &writeReturnRead, &closefd)) {
         return NULL;
     }
 
@@ -858,6 +859,7 @@ static ZstdCompressionWriter *ZstdCompressor_stream_writer(ZstdCompressor *self,
     result->bytesCompressed = 0;
     result->writeReturnRead =
         writeReturnRead ? PyObject_IsTrue(writeReturnRead) : 0;
+    result->closefd = closefd ? PyObject_IsTrue(closefd) : 1;
 
     return result;
 }

c-ext/python-zstandard.h

@@ -122,6 +122,7 @@ typedef struct {
     int entered;
     int closed;
     int writeReturnRead;
+    int closefd;
     unsigned long long bytesCompressed;
 } ZstdCompressionWriter;
 

docs/compressor.rst

@@ -269,6 +269,9 @@ were subsequently written to the compressor. ``True`` is the *proper* behavior
 for ``write()`` as specified by the ``io.RawIOBase`` interface and will become
 the default value in a future release.
 
+The ``closefd`` keyword argument defines whether to close the underlying stream
+when this instance is itself ``close()``d. The default is ``True``.
+
 Streaming Output API
 ====================
 

tests/test_compressor.py

@@ -950,6 +950,45 @@ def test_close(self):
             writer.write(b"foo")
 
         self.assertTrue(writer.closed)
+        self.assertTrue(buffer.closed)
+
+    def test_close_closefd_false(self):
+        buffer = NonClosingBytesIO()
+        cctx = zstd.ZstdCompressor(level=1)
+        writer = cctx.stream_writer(buffer, closefd=False)
+
+        writer.write(b"foo" * 1024)
+        self.assertFalse(writer.closed)
+        self.assertFalse(buffer.closed)
+        writer.close()
+        self.assertTrue(writer.closed)
+        self.assertFalse(buffer.closed)
+
+        with self.assertRaisesRegex(ValueError, "stream is closed"):
+            writer.write(b"foo")
+
+        with self.assertRaisesRegex(ValueError, "stream is closed"):
+            writer.flush()
+
+        with self.assertRaisesRegex(ValueError, "stream is closed"):
+            with writer:
+                pass
+
+        self.assertEqual(
+            buffer.getvalue(),
+            b"\x28\xb5\x2f\xfd\x00\x48\x55\x00\x00\x18\x66\x6f"
+            b"\x6f\x01\x00\xfa\xd3\x77\x43",
+        )
+
+        # Context manager exit should close stream.
+        buffer = io.BytesIO()
+        writer = cctx.stream_writer(buffer, closefd=False)
+
+        with writer:
+            writer.write(b"foo")
+
+        self.assertTrue(writer.closed)
+        self.assertFalse(buffer.closed)
 
     def test_empty(self):
         buffer = NonClosingBytesIO()

zstandard/__init__.pyi

@@ -296,6 +296,8 @@ class ZstdCompressor(object):
         size: int = -1,
         write_size: int = -1,
         write_return_read: bool = None,
+        *,
+        closefd: bool = True,
     ) -> ZstdCompressionWriter: ...
     def read_to_iter(
         self,

zstandard/backend_cffi.py

@@ -490,12 +490,19 @@ def _get_compression_parameter(params, param):
 
 class ZstdCompressionWriter(object):
     def __init__(
-        self, compressor, writer, source_size, write_size, write_return_read
+        self,
+        compressor,
+        writer,
+        source_size,
+        write_size,
+        write_return_read,
+        closefd=True,
     ):
         self._compressor = compressor
         self._writer = writer
         self._write_size = write_size
         self._write_return_read = bool(write_return_read)
+        self._closefd = bool(closefd)
         self._entered = False
         self._closed = False
         self._bytes_compressed = 0
@@ -553,7 +560,7 @@ def close(self):
 
         # Call close() on underlying stream as well.
         f = getattr(self._writer, "close", None)
-        if f:
+        if self._closefd and f:
             f()
 
     @property
@@ -1512,6 +1519,7 @@ def stream_writer(
         size=-1,
         write_size=COMPRESSION_RECOMMENDED_OUTPUT_SIZE,
         write_return_read=False,
+        closefd=True,
     ):
 
         if not hasattr(writer, "write"):
@@ -1523,7 +1531,7 @@ def stream_writer(
             size = lib.ZSTD_CONTENTSIZE_UNKNOWN
 
         return ZstdCompressionWriter(
-            self, writer, size, write_size, write_return_read
+            self, writer, size, write_size, write_return_read, closefd=closefd
         )
 
     write_to = stream_writer