Merge pull request fsspec#660 from limx0/localfile_compression

Add compression to LocalFileOpener

Author: martindurant

fsspec/implementations/local.py

@@ -7,6 +7,8 @@
 import tempfile
 
 from fsspec import AbstractFileSystem
+from fsspec.compression import compr
+from fsspec.core import get_compression
 from fsspec.utils import stringify_path
 
 
@@ -210,19 +212,25 @@ def make_path_posix(path, sep=os.sep):
 
 
 class LocalFileOpener(io.IOBase):
-    def __init__(self, path, mode, autocommit=True, fs=None, **kwargs):
+    def __init__(
+        self, path, mode, autocommit=True, fs=None, compression=None, **kwargs
+    ):
         self.path = path
         self.mode = mode
         self.fs = fs
         self.f = None
         self.autocommit = autocommit
+        self.compression = get_compression(path, compression)
         self.blocksize = io.DEFAULT_BUFFER_SIZE
         self._open()
 
     def _open(self):
         if self.f is None or self.f.closed:
             if self.autocommit or "w" not in self.mode:
                 self.f = open(self.path, mode=self.mode)
+                if self.compression:
+                    compress = compr[self.compression]
+                    self.f = compress(self.f, mode=self.mode)
             else:
                 # TODO: check if path is writable?
                 i, name = tempfile.mkstemp()

fsspec/implementations/tests/test_local.py

@@ -1,5 +1,6 @@
 from __future__ import absolute_import, division, print_function
 
+import bz2
 import gzip
 import os
 import os.path
@@ -662,3 +663,19 @@ def test_transaction(tmpdir):
         read_content = fp.read()
 
     assert content == read_content
+
+
+@pytest.mark.parametrize(
+    "opener, ext", [(bz2.open, ".bz2"), (gzip.open, ".gz"), (open, "")]
+)
+def test_infer_compression(tmpdir, opener, ext):
+    filename = str(tmpdir / f"test{ext}")
+    content = b"hello world"
+    with opener(filename, "wb") as fp:
+        fp.write(content)
+
+    fs = LocalFileSystem()
+    with fs.open(f"file://{filename}", "rb", compression="infer") as fp:
+        read_content = fp.read()
+
+    assert content == read_content