memfs2: make it non-cacheable, add tests, minor refactoring and optimizations (#73)

memfs2: make it non-cacheable, add tests, minor refactoring

Author: skshetry

.pre-commit-config.yaml

@@ -26,7 +26,7 @@ repos:
       - id: codespell
         args:
         - --ignore-words-list
-        - fo
+        - fo,cachable
   - repo: https://github.com/asottile/pyupgrade
     rev: v2.31.0
     hooks:

src/dvc_objects/fs/implementations/_memory.py

@@ -1,13 +1,15 @@
 import errno
+import os
 
 from fsspec import AbstractFileSystem
 from fsspec.implementations.memory import MemoryFile
 from fsspec.implementations.memory import MemoryFileSystem as MemFS
 
 
 class MemFS2(AbstractFileSystem):  # pylint: disable=abstract-method
-    """In-Memory Object Storage FileSystem."""
+    """In-Memory Object Storage FileSystem based on Trie data-structure."""
 
+    cachable = False
     protocol = "memory"
     root_marker = "/"
     _strip_protocol = MemFS._strip_protocol  # pylint: disable=protected-access
@@ -20,19 +22,14 @@ def __init__(self, *args, **storage_options):
 
     def ls(self, path, detail=False, **kwargs):
         path = self._strip_protocol(path)
-        if path not in ("", "/"):
-            info = self.info(path)
-            if info["type"] != "directory":
-                return [info] if detail else [path]
+        out = []
 
-        ret = []
-
-        def node_factory(path_conv, paths, children, *args):
-            node_path = path_conv(paths)
-            if path == node_path:
-                list(filter(None, children))
+        def node_factory(path_conv, parts, children, _filelike=None):
+            node_path = path_conv(parts)
+            if path == node_path and children:
+                list(children)
             else:
-                ret.append(node_path)
+                out.append(node_path)
 
         try:
             self.trie.traverse(node_factory, prefix=path)
@@ -44,9 +41,28 @@ def node_factory(path_conv, paths, children, *args):
             ) from exc
 
         if not detail:
-            return ret
+            return out
+        return [self.info(p) for p in out]
 
-        return [self.info(_path) for _path in ret]
+    def info(self, path, **kwargs):
+        path = self._strip_protocol(path)
+        if path in ("", "/") or self.trie.has_subtrie(path):
+            return {
+                "name": path,
+                "size": 0,
+                "type": "directory",
+            }
+
+        if filelike := self.trie.get(path):
+            return {
+                "name": path,
+                "size": filelike.size
+                if hasattr(filelike, "size")
+                else filelike.getbuffer().nbytes,
+                "type": "file",
+                "created": getattr(filelike, "created", None),
+            }
+        raise FileNotFoundError(errno.ENOENT, "No such file", path)
 
     def _rm(self, path):
         path = self._strip_protocol(path)
@@ -58,6 +74,11 @@ def _rm(self, path):
         except KeyError as e:
             raise FileNotFoundError(errno.ENOENT, "No such file", path) from e
 
+    def rm(self, path, recursive=False, maxdepth=None):
+        paths = self.expand_path(path, recursive=recursive, maxdepth=maxdepth)
+        for p in paths:
+            self.store.pop(p, None)
+
     def _open(  # pylint: disable=arguments-differ
         self, path, mode="rb", **kwargs
     ):
@@ -70,47 +91,26 @@ def _open(  # pylint: disable=arguments-differ
 
         if mode in ["rb", "ab", "rb+"]:
             if filelike := self.trie.get(path):
-                if mode == "ab":
-                    # position at the end of file
-                    filelike.seek(0, 2)
-                else:
-                    # position at the beginning of file
-                    filelike.seek(0)
+                filelike.seek(0, os.SEEK_END if mode == "ab" else os.SEEK_SET)
                 return filelike
             raise FileNotFoundError(errno.ENOENT, "No such file", path)
-        if mode == "wb":
-            filelike = MemoryFile(self, path)
-            if not self._intrans:
-                filelike.commit()
-            return filelike
 
-    def info(self, path, **kwargs):
-        path = self._strip_protocol(path)
-        if path in ("", "/") or self.trie.has_subtrie(path):
-            return {
-                "name": path,
-                "size": 0,
-                "type": "directory",
-            }
-
-        if filelike := self.trie.get(path):
-            return {
-                "name": path,
-                "size": filelike.size
-                if hasattr(filelike, "size")
-                else filelike.getbuffer().nbytes,
-                "type": "file",
-                "created": getattr(filelike, "created", None),
-            }
-        raise FileNotFoundError(errno.ENOENT, "No such file", path)
+        assert mode == "wb"
+        filelike = MemoryFile(self, path)
+        if not self._intrans:
+            filelike.commit()
+        return filelike
 
     def cp_file(self, path1, path2, **kwargs):
         path1 = self._strip_protocol(path1)
         path2 = self._strip_protocol(path2)
-        if self.isdir(path1):
+
+        try:
+            src = self.open(path1, "rb")
+        except IsADirectoryError:
             return
 
-        with self.open(path1, "rb") as src, self.open(path2, "wb") as dst:
+        with src, self.open(path2, "wb") as dst:
             dst.write(src.getbuffer())
 
     def created(self, path):

tests/fs/test_memfs.py

@@ -0,0 +1,211 @@
+import os
+from unittest.mock import ANY
+
+import pytest
+
+from dvc_objects.fs.implementations._memory import MemFS2
+
+
+@pytest.fixture
+def m():
+    return MemFS2()
+
+
+def test_memfs_should_not_be_cached():
+    assert MemFS2() is not MemFS2()
+
+
+def test_1(m):
+    m.touch("/somefile")  # NB: is found with or without initial /
+    m.touch("afiles/and/another")
+    files = m.find("")
+    assert files == ["/afiles/and/another", "/somefile"]
+
+    files = sorted(m.get_mapper())
+    assert files == ["afiles/and/another", "somefile"]
+
+
+def test_strip(m):
+    assert m._strip_protocol("") == ""
+    assert m._strip_protocol("memory://") == ""
+    assert m._strip_protocol("afile") == "/afile"
+    assert m._strip_protocol("/b/c") == "/b/c"
+    assert m._strip_protocol("/b/c/") == "/b/c"
+
+
+def test_put_single(m, tmpdir):
+    fn = os.path.join(str(tmpdir), "dir")
+    os.mkdir(fn)
+    open(os.path.join(fn, "abc"), "w").write("text")
+    m.put(fn, "/test")  # no-op, no files
+    assert not m.exists("/test/abc")
+    assert not m.exists("/test/dir")
+    m.put(fn + "/", "/test", recursive=True)
+    assert m.cat("/test/abc") == b"text"
+
+
+def test_ls(m):
+    m.touch("/dir/afile")
+    m.touch("/dir/dir1/bfile")
+    m.touch("/dir/dir1/cfile")
+
+    assert m.ls("/", False) == ["/dir"]
+    assert m.ls("/dir", False) == ["/dir/afile", "/dir/dir1"]
+    assert m.ls("/dir", True)[0]["type"] == "file"
+    assert m.ls("/dir", True)[1]["type"] == "directory"
+
+    assert len(m.ls("/dir/dir1")) == 2
+    assert m.ls("/dir/afile") == ["/dir/afile"]
+    assert m.ls("/dir/dir1/bfile") == ["/dir/dir1/bfile"]
+    assert m.ls("/dir/dir1/cfile") == ["/dir/dir1/cfile"]
+
+    with pytest.raises(FileNotFoundError):
+        m.ls("/dir/not-existing-file")
+
+
+def test_mv_recursive(m):
+    m.mkdir("src")
+    m.touch("src/file.txt")
+    m.mv("src", "dest", recursive=True)
+    assert m.exists("dest/file.txt")
+    assert not m.exists("src")
+
+
+def test_rm(m):
+    m.touch("/dir1/dir2/file")
+    m.rm("/dir1", recursive=True)
+    assert not m.exists("/dir1/dir2/file")
+    assert not m.exists("/dir1/dir2")
+    assert not m.exists("/dir1")
+
+    with pytest.raises(FileNotFoundError):
+        m.rm("/dir1", recursive=True)
+
+
+def test_rm_multiple_files(m):
+    m.touch("/dir/file1")
+    m.touch("/dir/file2")
+
+    m.rm(["/dir/file1", "/dir/file2"])
+    assert not m.ls("/")
+
+
+def test_rm_file(m):
+    m.touch("/dir/file")
+    with pytest.raises(IsADirectoryError):
+        m.rm_file("/dir")
+
+    with pytest.raises(FileNotFoundError):
+        m.rm_file("/dir/foo")
+
+    m.rm_file("/dir/file")
+    assert not m.exists("/dir/file")
+
+
+def test_rewind(m):
+    # https://github.com/fsspec/filesystem_spec/issues/349
+    with m.open("src/file.txt", "w") as f:
+        f.write("content")
+    with m.open("src/file.txt") as f:
+        assert f.tell() == 0
+
+
+def test_no_rewind_append_mode(m):
+    # https://github.com/fsspec/filesystem_spec/issues/349
+    with m.open("src/file.txt", "w") as f:
+        f.write("content")
+    with m.open("src/file.txt", "a") as f:
+        assert f.tell() == 7
+
+
+def test_seekable(m):
+    fn0 = "foo.txt"
+    with m.open(fn0, "wb") as f:
+        f.write(b"data")
+
+    f = m.open(fn0, "rt")
+    assert f.seekable(), "file is not seekable"
+    f.seek(1)
+    assert f.read(1) == "a"
+    assert f.tell() == 2
+
+
+def test_try_open_directory(m):
+    m.touch("/dir/file")
+    with pytest.raises(IsADirectoryError):
+        m.open("dir")
+
+
+def test_try_open_not_existing_file(m):
+    with pytest.raises(FileNotFoundError):
+        m.open("not-existing-file")
+
+
+def test_try_open_file_on_super_prefix(m):
+    m.touch("/afile")
+    with pytest.raises(NotADirectoryError):
+        m.open("/afile/file")
+
+
+def test_empty_raises(m):
+    with pytest.raises(FileNotFoundError):
+        m.ls("nonexistent")
+
+    with pytest.raises(FileNotFoundError):
+        m.info("nonexistent")
+
+
+def test_moves(m):
+    m.touch("source.txt")
+    m.mv("source.txt", "target.txt")
+
+    m.touch("source2.txt")
+    m.mv("source2.txt", "target2.txt", recursive=True)
+    assert m.find("") == ["/target.txt", "/target2.txt"]
+
+
+def test_remove_all(m: MemFS2):
+    m.touch("afile")
+    m.rm("/", recursive=True)
+    assert not m.ls("/")
+
+
+def test_created(m):
+    m.touch("/dir/afile")
+    assert m.created("/dir/afile") == m.trie["/dir/afile"].created
+    assert m.created("/dir") is None
+
+
+def test_info(m):
+    m.touch("/dir/file")
+
+    assert m.info("/") == {"name": "", "size": 0, "type": "directory"}
+    assert m.info("/dir") == {"name": "/dir", "size": 0, "type": "directory"}
+    assert m.info("/dir/file") == {
+        "name": "/dir/file",
+        "size": 0,
+        "type": "file",
+        "created": ANY,
+    }
+
+    with pytest.raises(FileNotFoundError):
+        m.info("/not-existing-file")
+
+
+def test_cp_file(m):
+    m.pipe_file("/afile", b"content")
+    m.cp_file("/afile", "/bfile")
+    assert m.cat_file("/bfile") == m.cat_file("/afile") == b"content"
+
+
+def test_transaction(m):
+    m.start_transaction()
+    m.touch("/dir/afile")
+    assert m.find("/") == []
+    m.end_transaction()
+    assert m.find("/") == ["/dir/afile"]
+
+    with m.transaction:
+        m.touch("/dir/bfile")
+        assert m.find("/") == ["/dir/afile"]
+    assert m.find("/") == ["/dir/afile", "/dir/bfile"]