From 778d5fd2d40959af0674340038e102e112a6dcab Mon Sep 17 00:00:00 2001 From: TWTom041 Date: Mon, 22 Jul 2024 23:38:03 +0800 Subject: [PATCH] implement memoryfs save() and load() --- fsspec/implementations/memory.py | 13 ++++++++++++- fsspec/implementations/tests/test_memory.py | 17 +++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/fsspec/implementations/memory.py b/fsspec/implementations/memory.py index 93860af6a..dc57cb99d 100644 --- a/fsspec/implementations/memory.py +++ b/fsspec/implementations/memory.py @@ -5,7 +5,7 @@ from errno import ENOTEMPTY from io import BytesIO from pathlib import PurePath, PureWindowsPath -from typing import Any, ClassVar +from typing import Any, ClassVar, BinaryIO from fsspec import AbstractFileSystem from fsspec.implementations.local import LocalFileSystem @@ -265,6 +265,17 @@ def rm(self, path, recursive=False, maxdepth=None): self.rm_file(p) else: self.rmdir(p) + + def save(self, file: BinaryIO): + import pickle + pickle.dump({"store": self.store, "pseudo_dirs": self.pseudo_dirs}, file) + + def load(self, file: BinaryIO): + import pickle + _internal = pickle.load(file) + self.store = _internal["store"] + self.pseudo_dirs = _internal["pseudo_dirs"] + class MemoryFile(BytesIO): diff --git a/fsspec/implementations/tests/test_memory.py b/fsspec/implementations/tests/test_memory.py index 600022a03..9ca7065b1 100644 --- a/fsspec/implementations/tests/test_memory.py +++ b/fsspec/implementations/tests/test_memory.py @@ -380,3 +380,20 @@ def test_open_path_windows(m): f.write(b"some\nlines\nof\ntext") assert m.read_text(path) == "some\nlines\nof\ntext" + +def test_save_and_load(m): + from io import BytesIO + internal = BytesIO() + with m.open("/foo/bar/file", "wb") as f: + f.write(b"some\nlines\nof\ntext") + with m.open("/foo/file", "wb") as f: + f.write(b"some\nlines\nof\ntext") + m.makedirs("/foo/foo") + m.save(internal) + m.store = None + m.pseudo_dirs = None + internal.seek(0) + m.load(internal) + assert m.cat("/foo/file") == b"some\nlines\nof\ntext" + assert m.ls("/", False) == ["/foo"] + assert sorted(m.ls("/foo", False)) == sorted(["/foo/foo", "/foo/bar", "/foo/file"])