Memfs no bytes copy (#999)

* memory: Avoid bytes copies where possible

* Pass 'data' kwarg into MemoryFile from open

* Implement pipe_file that initializes rather than copies

* Follow CI linter recommendations

* remove assert and unused kwargs

* Don't copy bytes in MemoryFS

* remove test

Co-authored-by: Brian Ealdwine <[email protected]>

Author: martindurant

fsspec/implementations/memory.py

@@ -38,7 +38,7 @@ def ls(self, path, detail=False, **kwargs):
             return [
                 {
                     "name": path,
-                    "size": self.store[path].getbuffer().nbytes,
+                    "size": self.store[path].size,
                     "type": "file",
                     "created": self.store[path].created,
                 }
@@ -53,7 +53,7 @@ def ls(self, path, detail=False, **kwargs):
                     out.append(
                         {
                             "name": p2,
-                            "size": self.store[p2].getbuffer().nbytes,
+                            "size": self.store[p2].size,
                             "type": "file",
                             "created": self.store[p2].created,
                         }
@@ -114,6 +114,13 @@ def makedirs(self, path, exist_ok=False):
             if not exist_ok:
                 raise
 
+    def pipe_file(self, path, value, **kwargs):
+        """Set the bytes of given file
+
+        Avoids copies of the data if possible
+        """
+        self.open(path, "wb", data=value)
+
     def rmdir(self, path):
         path = self._strip_protocol(path)
         if path == "":
@@ -145,9 +152,7 @@ def info(self, path, **kwargs):
             filelike = self.store[path]
             return {
                 "name": path,
-                "size": filelike.size
-                if hasattr(filelike, "size")
-                else filelike.getbuffer().nbytes,
+                "size": filelike.size,
                 "type": "file",
                 "created": getattr(filelike, "created", None),
             }
@@ -184,7 +189,7 @@ def _open(
             else:
                 raise FileNotFoundError(path)
         if mode == "wb":
-            m = MemoryFile(self, path)
+            m = MemoryFile(self, path, kwargs.get("data"))
             if not self._intrans:
                 m.commit()
             return m
@@ -193,7 +198,9 @@ def cp_file(self, path1, path2, **kwargs):
         path1 = self._strip_protocol(path1)
         path2 = self._strip_protocol(path2)
         if self.isfile(path1):
-            self.store[path2] = MemoryFile(self, path2, self.store[path1].getbuffer())
+            self.store[path2] = MemoryFile(
+                self, path2, self.store[path1].getvalue()
+            )  # implicit copy
         elif self.isdir(path1):
             if path2 not in self.pseudo_dirs:
                 self.pseudo_dirs.append(path2)
@@ -203,7 +210,7 @@ def cp_file(self, path1, path2, **kwargs):
     def cat_file(self, path, start=None, end=None, **kwargs):
         path = self._strip_protocol(path)
         try:
-            return self.store[path].getvalue()[start:end]
+            return bytes(self.store[path].getbuffer()[start:end])
         except KeyError:
             raise FileNotFoundError(path)
 
@@ -246,17 +253,18 @@ def __init__(self, fs=None, path=None, data=None):
         self.path = path
         self.created = datetime.utcnow().timestamp()
         if data:
-            self.write(data)
-            self.size = len(data)
+            super().__init__(data)
             self.seek(0)
 
+    @property
+    def size(self):
+        return self.getbuffer().nbytes
+
     def __enter__(self):
         return self
 
     def close(self):
-        position = self.tell()
-        self.size = self.seek(0, 2)
-        self.seek(position)
+        pass
 
     def discard(self):
         pass

fsspec/tests/test_utils.py

@@ -4,6 +4,7 @@
 
 import pytest
 
+import fsspec.utils
 from fsspec.utils import (
     can_be_local,
     common_prefix,
@@ -377,3 +378,9 @@ def test_merge_offset_ranges(max_gap, max_block):
     assert expect_paths == result_paths
     assert expect_starts == result_starts
     assert expect_ends == result_ends
+
+
+def test_size():
+    f = io.BytesIO(b"hello")
+    assert fsspec.utils.file_size(f) == 5
+    assert f.tell() == 0

fsspec/utils.py

@@ -619,3 +619,12 @@ def __exit__(
     # forward anything else too
     def __getattr__(self, name):
         return getattr(self.fp, name)
+
+
+def file_size(filelike):
+    """Find length of any open read-mode file-like"""
+    pos = filelike.tell()
+    try:
+        return filelike.seek(0, 2)
+    finally:
+        filelike.seek(pos)