Close underlying files when using compression

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)
 

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}

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):

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"