diff --git a/NEWS.rst b/NEWS.rst index 00c3ce1c..96492586 100644 --- a/NEWS.rst +++ b/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 diff --git a/README.rst b/README.rst index 8fb9ff8c..bf877ec4 100644 --- a/README.rst +++ b/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*. diff --git a/c-ext/decompressionreader.c b/c-ext/decompressionreader.c index b38502ef..523d1d12 100644 --- a/c-ext/decompressionreader.c +++ b/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; diff --git a/c-ext/decompressor.c b/c-ext/decompressor.c index c8878999..be10609d 100644 --- a/c-ext/decompressor.c +++ b/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; diff --git a/tests/test_decompressor.py b/tests/test_decompressor.py index add8dc41..6668ec83 100644 --- a/tests/test_decompressor.py +++ b/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): diff --git a/zstd_cffi.py b/zstd_cffi.py index 494ba502..a43c9afe 100644 --- a/zstd_cffi.py +++ b/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')