diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 42ed4be42..775b207c4 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -6,6 +6,9 @@ Dev Fixes +- Close the underlying file when closing a stream opened with compression + through ``AbstractFileSystem.open``, and propagate errors raised on close (#1672) + - Avoid mutating live ``BlockCache`` and ``BackgroundBlockCache`` instances when pickling (#2102) diff --git a/fsspec/compression.py b/fsspec/compression.py index 11c2e3d3f..3c69b8270 100644 --- a/fsspec/compression.py +++ b/fsspec/compression.py @@ -11,6 +11,43 @@ def noop_file(file, mode, **kwargs): return file +class _ClosingFile: + """A compression stream that owns its underlying file.""" + + def __init__(self, file, raw): + self._file = file + self._raw = raw + + def __getattr__(self, name): + return getattr(self._file, name) + + def __iter__(self): + return self + + def __next__(self): + return next(self._file) + + def __enter__(self): + self._file.__enter__() + return self + + def __exit__(self, *args): + self.close() + + def close(self): + try: + # Finalize the compression stream before flushing the raw file. + if not self._file.closed: + self._file.close() + finally: + # Some codecs already close their input, while others leave it open. + if not self._raw.closed: + self._raw.close() + + def __del__(self): + self.close() + + # TODO: files should also be available as contexts # should be functions of the form func(infile, mode=, **kwargs) -> file-like compr = {None: noop_file} diff --git a/fsspec/spec.py b/fsspec/spec.py index 0e1908f52..a90dc0159 100644 --- a/fsspec/spec.py +++ b/fsspec/spec.py @@ -1417,12 +1417,14 @@ def open( if not ac and "r" not in mode: self.transaction.files.append(f) if compression is not None: - from fsspec.compression import compr + from fsspec.compression import _ClosingFile, compr from fsspec.core import get_compression compression = get_compression(path, compression) compress = compr[compression] - f = compress(f, mode=mode[0]) + compressed = compress(f, mode=mode[0]) + if compressed is not f: + f = _ClosingFile(compressed, f) return f def touch(self, path, truncate=True, **kwargs): diff --git a/fsspec/tests/test_compression.py b/fsspec/tests/test_compression.py index 462b35904..b52644691 100644 --- a/fsspec/tests/test_compression.py +++ b/fsspec/tests/test_compression.py @@ -1,3 +1,4 @@ +import io import pathlib import sys @@ -8,6 +9,176 @@ from fsspec.utils import compressions, infer_compression +@pytest.mark.parametrize("compression", ["gzip", "bz2", "lzma"]) +@pytest.mark.parametrize("mode", ["rb", "rt", "wb", "wt"]) +@pytest.mark.parametrize("use_context", [False, True]) +def test_fs_open_closes_raw_file(compression, mode, use_context, monkeypatch): + codec = pytest.importorskip(compression) + data = b"hello\nworld\n" + raw = io.BytesIO(codec.compress(data) if "r" in mode else b"") + fs = fsspec.AbstractFileSystem() + monkeypatch.setattr(fs, "_open", lambda *args, **kwargs: raw) + + f = fs.open("test", mode, compression=compression) + try: + value = data if "b" in mode else data.decode() + if "r" in mode: + assert f.read() == value + else: + f.write(value) + if use_context: + with f: + pass + else: + f.close() + assert f.closed + assert raw.closed + f.close() + finally: + f.close() + raw.close() + + +@pytest.mark.parametrize("compression", ["gzip", "bz2", "lzma"]) +@pytest.mark.parametrize("mode", ["wb", "wt"]) +def test_fs_open_propagates_raw_close_error(compression, mode, monkeypatch): + pytest.importorskip(compression) + + class FailingCloseFile(io.BytesIO): + def close(self): + if not self.closed: + super().close() + raise OSError("upload failed on close") + + raw = FailingCloseFile() + fs = fsspec.AbstractFileSystem() + monkeypatch.setattr(fs, "_open", lambda *args, **kwargs: raw) + + try: + with pytest.raises(OSError, match="upload failed on close"): + with fs.open("test", mode, compression=compression) as f: + f.write(b"data" if "b" in mode else "data") + assert raw.closed + finally: + # Suppress the expected error while cleaning up on the unfixed version. + if not raw.closed: + with pytest.raises(OSError, match="upload failed on close"): + raw.close() + + +@pytest.mark.parametrize("compression", ["gzip", "bz2", "lzma"]) +def test_fs_open_closes_raw_file_on_compression_error(compression, monkeypatch): + pytest.importorskip(compression) + raw = io.BytesIO() + fs = fsspec.AbstractFileSystem() + monkeypatch.setattr(fs, "_open", lambda *args, **kwargs: raw) + f = fs.open("test", "wb", compression=compression) + f.write(b"data") + + def fail_write(data): + raise OSError("writing compression trailer failed") + + monkeypatch.setattr(raw, "write", fail_write) + try: + with pytest.raises(OSError, match="writing compression trailer failed"): + f.close() + assert raw.closed + finally: + raw.close() + + +@pytest.mark.parametrize("compression", ["gzip", "bz2", "lzma"]) +@pytest.mark.parametrize("commit", [False, True]) +def test_fs_open_compression_transaction(compression, commit, tmp_path): + codec = pytest.importorskip(compression) + fs = fsspec.filesystem("file") + path = tmp_path / "test" + fs.start_transaction() + transaction = fs.transaction + try: + with fs.open(path, "wt", compression=compression, newline="\n") as f: + f.write("hello\nworld\n") + assert not path.exists() + assert len(transaction.files) == 1 + assert transaction.files[0].closed + finally: + transaction.complete(commit=commit) + assert path.exists() == commit + if commit: + assert codec.decompress(path.read_bytes()) == b"hello\nworld\n" + + +def test_fs_open_infer_no_compression(monkeypatch): + raw = io.BytesIO(b"data") + fs = fsspec.AbstractFileSystem() + monkeypatch.setattr(fs, "_open", lambda *args, **kwargs: raw) + with fs.open("test.unknown", compression="infer") as f: + assert f is raw + assert f.read() == b"data" + + +@pytest.mark.parametrize("compression", ["gzip", "bz2", "lzma"]) +def test_fs_open_compression_file_methods(compression, monkeypatch): + codec = pytest.importorskip(compression) + data = b"hello\nworld\n" + raw = io.BytesIO(codec.compress(data)) + fs = fsspec.AbstractFileSystem() + monkeypatch.setattr(fs, "_open", lambda *args, **kwargs: raw) + with fs.open("test", "rb", compression=compression) as f: + assert list(f) == [b"hello\n", b"world\n"] + f.seek(0) + buffer = bytearray(len(data)) + assert f.readinto(buffer) == len(data) + assert buffer == data + with pytest.raises(ValueError): + with f: + pass + + +@pytest.mark.parametrize("compression", list(compr)) +@pytest.mark.parametrize("mode", ["b", "t"]) +def test_fs_open_compression_roundtrip(compression, mode, tmp_path): + fs = fsspec.filesystem("file") + path = tmp_path / "test" + data = b"hello\nworld\n" if mode == "b" else "hello\nworld\n" + with fs.open(path, "w" + mode, compression=compression) as f: + f.write(data) + with fs.open(path, "r" + mode, compression=compression) as f: + assert f.read() == data + + +@pytest.mark.parametrize("mode", ["rb", "rt", "wb", "wt"]) +def test_fs_open_zstandard_closes_raw_once(mode, monkeypatch): + zstd = pytest.importorskip("zstandard") + + def compress(raw, mode): + if mode == "r": + return zstd.ZstdDecompressor().stream_reader(raw) + return zstd.ZstdCompressor().stream_writer(raw) + + class CountingFile(io.BytesIO): + close_calls = 0 + + def close(self): + self.close_calls += 1 + super().close() + + data = b"data" + raw = CountingFile(zstd.ZstdCompressor().compress(data) if "r" in mode else b"") + fs = fsspec.AbstractFileSystem() + monkeypatch.setattr(fs, "_open", lambda *args, **kwargs: raw) + monkeypatch.setitem(compr, "zstd", compress) + with fs.open("test", mode, compression="zstd") as f: + value = data if "b" in mode else data.decode() + if "r" in mode: + assert f.read() == value + else: + f.write(value) + f.close() + assert raw.closed + assert raw.close_calls == 1 + + def test_infer_custom_compression(): """Inferred compression gets values from fsspec.compression.compr.""" assert infer_compression("fn.zip") == "zip"