decompressionreader: don't require usage within context manager (#34)

NEWS.rst

@@ -69,10 +69,15 @@ Backwards Compatibility Notes
   potentially expensive setup when using dictionaries. We now call
   ``ZSTD_CCtx_reset()`` on every operation and don't attempt to change
   compression parameters.
+* Objects returned by ``ZstdDecompressor.stream_reader()`` no longer needs to be
+  used as a context manager. The context manager interface still exists and its
+  behavior is unchanged.
 
 New Features
 ------------
 
+* ``ZstdDecompressor.stream_reader()`` no longer needs to be used as a context
+  manager (#34).
 * Bundled zstandard library upgraded from 1.3.4 to 1.3.6.
 
 Changes

README.rst

@@ -542,17 +542,24 @@ Stream Reader API
 
    with open(path, 'rb') as fh:
        dctx = zstd.ZstdDecompressor()
-       with dctx.stream_reader(fh) as reader:
-           while True:
-               chunk = reader.read(16384)
-               if not chunk:
-                   break
+       reader = dctx.stream_reader(fh)
+       while True:
+           chunk = reader.read(16384)
+            if not chunk:
+                break
 
-               # Do something with decompressed chunk.
+            # Do something with decompressed chunk.
 
-The stream can only be read within a context manager. When the context
-manager exits, the stream is closed and the underlying resource is
-released and future operations against the stream will fail.
+The stream can also be used as a context manager::
+
+   with open(path, 'rb') as fh:
+       dctx = zstd.ZstdDecompressor()
+       with dctx.stream_reader(fh) as reader:
+           ...
+
+When used as a context manager, the stream is closed and the underlying
+resources are released when the context manager exits. Future operations against
+the stream will fail.
 
 The ``source`` argument to ``stream_reader()`` can be any object with a
 ``read(size)`` method or any object implementing the *buffer protocol*.

c-ext/decompressionreader.c

@@ -47,10 +47,6 @@ static ZstdDecompressionReader* reader_enter(ZstdDecompressionReader* self) {
 		return NULL;
 	}
 
-	if (ensure_dctx(self->decompressor, 1)) {
-		return NULL;
-	}
-
 	self->entered = 1;
 
 	Py_INCREF(self);
@@ -128,11 +124,6 @@ static PyObject* reader_read(ZstdDecompressionReader* self, PyObject* args, PyOb
 	ZSTD_outBuffer output;
 	size_t zresult;
 
-	if (!self->entered) {
-		PyErr_SetString(ZstdError, "read() must be called from an active context manager");
-		return NULL;
-	}
-
 	if (self->closed) {
 		PyErr_SetString(PyExc_ValueError, "stream is closed");
 		return NULL;
@@ -281,11 +272,6 @@ static PyObject* reader_seek(ZstdDecompressionReader* self, PyObject* args) {
 	unsigned long long readAmount = 0;
 	size_t defaultOutSize = ZSTD_DStreamOutSize();
 
-	if (!self->entered) {
-		PyErr_SetString(ZstdError, "seek() must be called from an active context manager");
-		return NULL;
-	}
-
 	if (self->closed) {
 		PyErr_SetString(PyExc_ValueError, "stream is closed");
 		return NULL;

c-ext/decompressor.c

@@ -575,6 +575,10 @@ static ZstdDecompressionReader* Decompressor_stream_reader(ZstdDecompressor* sel
 		return NULL;
 	}
 
+	if (ensure_dctx(self, 1)) {
+		return NULL;
+	}
+
 	result = (ZstdDecompressionReader*)PyObject_CallObject((PyObject*)&ZstdDecompressionReaderType, NULL);
 	if (NULL == result) {
 		return NULL;

tests/test_decompressor.py

@@ -293,10 +293,6 @@ class TestDecompressor_stream_reader(unittest.TestCase):
     def test_context_manager(self):
         dctx = zstd.ZstdDecompressor()
 
-        reader = dctx.stream_reader(b'foo')
-        with self.assertRaisesRegexp(zstd.ZstdError, 'read\(\) must be called from an active'):
-            reader.read(1)
-
         with dctx.stream_reader(b'foo') as reader:
             with self.assertRaisesRegexp(ValueError, 'cannot __enter__ multiple times'):
                 with reader as reader2:
@@ -440,7 +436,7 @@ def test_read_after_exit(self):
             while reader.read(16):
                 pass
 
-        with self.assertRaisesRegexp(zstd.ZstdError, 'read\(\) must be called from an active'):
+        with self.assertRaisesRegexp(ValueError, 'stream is closed'):
             reader.read(10)
 
     def test_illegal_seeks(self):
@@ -474,8 +470,7 @@ def test_illegal_seeks(self):
             with self.assertRaisesRegexp(ValueError, 'stream is closed'):
                 reader.seek(4, os.SEEK_SET)
 
-        with self.assertRaisesRegexp(
-            zstd.ZstdError, 'seek\(\) must be called from an active context'):
+        with self.assertRaisesRegexp(ValueError, 'stream is closed'):
             reader.seek(0)
 
     def test_seek(self):
@@ -492,6 +487,22 @@ def test_seek(self):
             reader.seek(4, os.SEEK_CUR)
             self.assertEqual(reader.read(2), b'ar')
 
+    def test_no_context_manager(self):
+        source = b'foobar' * 60
+        cctx = zstd.ZstdCompressor()
+        frame = cctx.compress(source)
+
+        dctx = zstd.ZstdDecompressor()
+        reader = dctx.stream_reader(frame)
+
+        self.assertEqual(reader.read(6), b'foobar')
+        self.assertEqual(reader.read(18), b'foobar' * 3)
+
+        # Calling close prevents subsequent use.
+        reader.close()
+
+        with self.assertRaisesRegexp(ValueError, 'stream is closed'):
+            reader.read(6)
 
 @make_cffi
 class TestDecompressor_decompressobj(unittest.TestCase):

zstd_cffi.py

@@ -1291,8 +1291,6 @@ def __enter__(self):
         if self._entered:
             raise ValueError('cannot __enter__ multiple times')
 
-        self._decompressor._ensure_dctx()
-
         self._entered = True
         return self
 
@@ -1353,9 +1351,6 @@ def __next__(self):
     next = __next__
 
     def read(self, size=-1):
-        if not self._entered:
-            raise ZstdError('read() must be called from an active context manager')
-
         if self._closed:
             raise ValueError('stream is closed')
 
@@ -1430,10 +1425,6 @@ def get_input():
         return ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
 
     def seek(self, pos, whence=os.SEEK_SET):
-        if not self._entered:
-            raise ZstdError('seek() must be called from an active context '
-                            'manager')
-
         if self._closed:
             raise ValueError('stream is closed')