Repository: fsspec/filesystem_spec
Pull request #1612 (2024-05-28 13:58:05)

Extend JSON serialization capabilities
#247 first introduced JSON serialization of filesystem objects. However, it fails to handle `storage_args` and `storage_options` that contain `Path` or other filesystems. This can occur in a few cases, such as:

- When passing a `Path` instance to construct a `DirFileSystem` (not officially documented but the `Path` object is automatically converted to a string in the initializer, so it is functional)
- When passing a filesystem to wrapper filesystems such as `CachingFileSystem`, `DirFileSystem` and `ReferenceFileSystem`.

This PR makes it possible to serialize such filesystems.


DarkLight1337: > The whole encoder/decoder classes are quite long now, there's no reason they need to be in the main `spec` module - can they be pulled into a different file?

Sure - I'll move it to `fsspec.json`.

Patch:
diff --git a/README.md b/README.md
index 426bf1f..b65b791 100644
--- a/README.md
+++ b/README.md
@@ -51,13 +51,13 @@ CI runtime. For local use, pick a version suitable for you.
 mamba create -n fsspec -c conda-forge  python=3.9 -y
 conda activate fsspec
 
-# Standard dev test install.
-pip install -e ".[dev,test]"
+# Standard dev install with docs and tests.
+pip install -e ".[dev,doc,test]"
 
 # Full tests except for downstream
 pip install s3fs
 pip uninstall s3fs
-pip install -e .[dev,test_full]
+pip install -e .[dev,doc,test_full]
 pip install s3fs --no-deps
 pytest -v
 
diff --git a/docs/environment.yml b/docs/environment.yml
index 2ab1379..fe05bd7 100644
--- a/docs/environment.yml
+++ b/docs/environment.yml
@@ -3,9 +3,3 @@ channels:
   - defaults
 dependencies:
   - python=3.9
-  - docutils<0.17
-  - numpydoc
-  - sphinx_rtd_theme
-  - yarl
-  - pip:
-    - sphinx-design
diff --git a/docs/source/developer.rst b/docs/source/developer.rst
index 07777f1..c74dca9 100644
--- a/docs/source/developer.rst
+++ b/docs/source/developer.rst
@@ -107,9 +107,9 @@ The following can be used to install ``fsspec`` in development mode
 
    git clone https://github.com/fsspec/filesystem_spec
    cd filesystem_spec
-   pip install -e .
+   pip install -e .[dev,doc,test]
 
-A number of additional dependencies are required to run tests, see "ci/environment*.yml", as
+A number of additional dependencies are required to run tests in full, see "ci/environment*.yml", as
 well as Docker. Most implementation-specific tests should skip if their requirements are
 not met.
 
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 1e103fd..8eb6b50 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -69,7 +69,7 @@ The following libraries use ``fsspec`` internally for path and file handling:
 .. _xarray: https://docs.xarray.dev/
 .. _zarr: https://zarr.readthedocs.io/
 .. _DVC: https://dvc.org/
-.. _kedro: https://kedro.readthedocs.io/en/stable/01_introduction/01_introduction.html
+.. _kedro: https://docs.kedro.org/en/stable/tutorial/set_up_data.html#supported-data-locations
 .. _pyxet: https://github.com/xetdata/pyxet
 .. _Huggingface🤗 Datasets: https://github.com/huggingface/datasets
 .. _pyarrow: https://arrow.apache.org/docs/python/
diff --git a/fsspec/implementations/cached.py b/fsspec/implementations/cached.py
index 5c6c145..2d67052 100644
--- a/fsspec/implementations/cached.py
+++ b/fsspec/implementations/cached.py
@@ -434,6 +434,7 @@ class CachingFileSystem(AbstractFileSystem):
             "__hash__",
             "__eq__",
             "to_json",
+            "to_dict",
             "cache_size",
             "pipe_file",
             "pipe",
@@ -506,15 +507,6 @@ class CachingFileSystem(AbstractFileSystem):
             ^ hash(self.target_protocol)
         )
 
-    def to_json(self):
-        """Calculate JSON representation.
-
-        Not implemented yet for CachingFileSystem.
-        """
-        raise NotImplementedError(
-            "CachingFileSystem JSON representation not implemented"
-        )
-
 
 class WholeFileCacheFileSystem(CachingFileSystem):
     """Caches whole remote files on first access
diff --git a/fsspec/implementations/reference.py b/fsspec/implementations/reference.py
index 06254fe..6e3de91 100644
--- a/fsspec/implementations/reference.py
+++ b/fsspec/implementations/reference.py
@@ -956,6 +956,8 @@ class ReferenceFileSystem(AsyncFileSystem):
                 if v.startswith("base64:"):
                     self.references[k] = base64.b64decode(v[7:])
                 self.references[k] = v
+            elif isinstance(v, dict):
+                self.references[k] = json.dumps(v)
             elif self.templates:
                 u = v[0]
                 if "{{" in u:
@@ -968,8 +970,6 @@ class ReferenceFileSystem(AsyncFileSystem):
                     else:
                         u = _render_jinja(u)
                 self.references[k] = [u] if len(v) == 1 else [u, v[1], v[2]]
-            elif isinstance(v, dict):
-                self.references[k] = json.dumps(v)
             else:
                 self.references[k] = v
         self.references.update(self._process_gen(references.get("gen", [])))
diff --git a/fsspec/json.py b/fsspec/json.py
new file mode 100644
index 0000000..3d1e473
--- /dev/null
+++ b/fsspec/json.py
@@ -0,0 +1,81 @@
+import json
+from contextlib import suppress
+from pathlib import PurePath
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+from .registry import _import_class, get_filesystem_class
+from .spec import AbstractFileSystem
+
+
+class FilesystemJSONEncoder(json.JSONEncoder):
+    def default(self, o: Any) -> Any:
+        if isinstance(o, AbstractFileSystem):
+            return o.to_dict()
+        if isinstance(o, PurePath):
+            cls = type(o)
+            return {"cls": f"{cls.__module__}.{cls.__name__}", "str": str(o)}
+
+        return super().default(o)
+
+
+class FilesystemJSONDecoder(json.JSONDecoder):
+    def __init__(
+        self,
+        *,
+        object_hook: Optional[Callable[[Dict[str, Any]], Any]] = None,
+        parse_float: Optional[Callable[[str], Any]] = None,
+        parse_int: Optional[Callable[[str], Any]] = None,
+        parse_constant: Optional[Callable[[str], Any]] = None,
+        strict: bool = True,
+        object_pairs_hook: Optional[Callable[[List[Tuple[str, Any]]], Any]] = None,
+    ) -> None:
+        self.original_object_hook = object_hook
+
+        super().__init__(
+            object_hook=self.custom_object_hook,
+            parse_float=parse_float,
+            parse_int=parse_int,
+            parse_constant=parse_constant,
+            strict=strict,
+            object_pairs_hook=object_pairs_hook,
+        )
+
+    @classmethod
+    def try_resolve_path_cls(cls, dct: Dict[str, Any]):
+        with suppress(Exception):
+            fqp = dct["cls"]
+
+            path_cls = _import_class(fqp)
+
+            if issubclass(path_cls, PurePath):
+                return path_cls
+
+        return None
+
+    @classmethod
+    def try_resolve_fs_cls(cls, dct: Dict[str, Any]):
+        with suppress(Exception):
+            if "cls" in dct:
+                try:
+                    fs_cls = _import_class(dct["cls"])
+                    if issubclass(fs_cls, AbstractFileSystem):
+                        return fs_cls
+                except Exception:
+                    if "protocol" in dct:  # Fallback if cls cannot be imported
+                        return get_filesystem_class(dct["protocol"])
+
+                    raise
+
+        return None
+
+    def custom_object_hook(self, dct: Dict[str, Any]):
+        if "cls" in dct:
+            if (obj_cls := self.try_resolve_fs_cls(dct)) is not None:
+                return AbstractFileSystem.from_dict(dct)
+            if (obj_cls := self.try_resolve_path_cls(dct)) is not None:
+                return obj_cls(dct["str"])
+
+        if self.original_object_hook is not None:
+            return self.original_object_hook(dct)
+
+        return dct
diff --git a/fsspec/registry.py b/fsspec/registry.py
index 85504d6..c261b9b 100644
--- a/fsspec/registry.py
+++ b/fsspec/registry.py
@@ -257,27 +257,33 @@ update the current installation.
 """
 
 
-def _import_class(cls, minv=None):
-    """Take a string FQP and return the imported class or identifier
+def _import_class(fqp: str):
+    """Take a fully-qualified path and return the imported class or identifier.
 
-    cls is of the form "package.module.klass" or "package.module:subobject.klass"
+    ``fqp`` is of the form "package.module.klass" or
+    "package.module:subobject.klass".
+
+    Warnings
+    --------
+    This can import arbitrary modules. Make sure you haven't installed any modules
+    that may execute malicious code at import time.
     """
-    if ":" in cls:
-        mod, name = cls.rsplit(":", 1)
-        s3 = mod == "s3fs"
-        mod = importlib.import_module(mod)
-        if s3 and mod.__version__.split(".") < ["0", "5"]:
-            warnings.warn(s3_msg)
-        for part in name.split("."):
-            mod = getattr(mod, part)
-        return mod
+    if ":" in fqp:
+        mod, name = fqp.rsplit(":", 1)
     else:
-        mod, name = cls.rsplit(".", 1)
-        s3 = mod == "s3fs"
-        mod = importlib.import_module(mod)
-        if s3 and mod.__version__.split(".") < ["0", "5"]:
-            warnings.warn(s3_msg)
-        return getattr(mod, name)
+        mod, name = fqp.rsplit(".", 1)
+
+    is_s3 = mod == "s3fs"
+    mod = importlib.import_module(mod)
+    if is_s3 and mod.__version__.split(".") < ["0", "5"]:
+        warnings.warn(s3_msg)
+    for part in name.split("."):
+        mod = getattr(mod, part)
+
+    if not isinstance(mod, type):
+        raise TypeError(f"{fqp} is not a class")
+
+    return mod
 
 
 def filesystem(protocol, **storage_options):
diff --git a/fsspec/spec.py b/fsspec/spec.py
index bcc01f5..f55a536 100644
--- a/fsspec/spec.py
+++ b/fsspec/spec.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import io
+import json
 import logging
 import os
 import threading
@@ -9,7 +10,7 @@ import weakref
 from errno import ESPIPE
 from glob import has_magic
 from hashlib import sha256
-from typing import ClassVar
+from typing import Any, ClassVar, Dict, Tuple
 
 from .callbacks import DEFAULT_CALLBACK
 from .config import apply_config, conf
@@ -115,6 +116,10 @@ class AbstractFileSystem(metaclass=_Cached):
     #: Extra *class attributes* that should be considered when hashing.
     _extra_tokenize_attributes = ()
 
+    # Set by _Cached metaclass
+    storage_args: Tuple[Any, ...]
+    storage_options: Dict[str, Any]
+
     def __init__(self, *args, **storage_options):
         """Create and configure file-system instance
 
@@ -1381,61 +1386,98 @@ class AbstractFileSystem(metaclass=_Cached):
                 length = size - offset
             return read_block(f, offset, length, delimiter)
 
-    def to_json(self):
+    def to_json(self) -> str:
         """
-        JSON representation of this filesystem instance
+        JSON representation of this filesystem instance.
 
         Returns
         -------
-        str: JSON structure with keys cls (the python location of this class),
-            protocol (text name of this class's protocol, first one in case of
-            multiple), args (positional args, usually empty), and all other
-            kwargs as their own keys.
+        JSON string with keys ``cls`` (the python location of this class),
+        protocol (text name of this class's protocol, first one in case of
+        multiple), ``args`` (positional args, usually empty), and all other
+        keyword arguments as their own keys.
+        """
+        from .json import FilesystemJSONEncoder
+
+        return json.dumps(self, cls=FilesystemJSONEncoder)
+
+    @staticmethod
+    def from_json(blob: str) -> AbstractFileSystem:
         """
-        import json
+        Recreate a filesystem instance from JSON representation.
+
+        See ``.to_json()`` for the expected structure of the input.
+
+        Parameters
+        ----------
+        blob: str
+
+        Returns
+        -------
+        file system instance, not necessarily of this particular class.
 
+        Warnings
+        --------
+        This can import arbitrary modules (as determined by the ``cls`` key).
+        Make sure you haven't installed any modules that may execute malicious code
+        at import time.
+        """
+        from .json import FilesystemJSONDecoder
+
+        return json.loads(blob, cls=FilesystemJSONDecoder)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        JSON-serializable dictionary representation of this filesystem instance.
+
+        Returns
+        -------
+        Dictionary with keys ``cls`` (the python location of this class),
+        protocol (text name of this class's protocol, first one in case of
+        multiple), ``args`` (positional args, usually empty), and all other
+        keyword arguments as their own keys.
+        """
         cls = type(self)
-        cls = ".".join((cls.__module__, cls.__name__))
-        proto = (
-            self.protocol[0]
-            if isinstance(self.protocol, (tuple, list))
-            else self.protocol
-        )
-        return json.dumps(
-            dict(
-                cls=cls,
-                protocol=proto,
-                args=self.storage_args,
-                **self.storage_options,
-            )
+        proto = self.protocol
+
+        return dict(
+            cls=f"{cls.__module__}:{cls.__name__}",
+            protocol=proto[0] if isinstance(proto, (tuple, list)) else proto,
+            args=self.storage_args,
+            **self.storage_options,
         )
 
     @staticmethod
-    def from_json(blob):
+    def from_dict(dct: Dict[str, Any]) -> AbstractFileSystem:
         """
-        Recreate a filesystem instance from JSON representation
+        Recreate a filesystem instance from dictionary representation.
 
-        See ``.to_json()`` for the expected structure of the input
+        See ``.to_dict()`` for the expected structure of the input.
 
         Parameters
         ----------
-        blob: str
+        dct: Dict[str, Any]
 
         Returns
         -------
         file system instance, not necessarily of this particular class.
+
+        Warnings
+        --------
+        This can import arbitrary modules (as determined by the ``cls`` key).
+        Make sure you haven't installed any modules that may execute malicious code
+        at import time.
         """
-        import json
+        from .json import FilesystemJSONDecoder
 
-        from .registry import _import_class, get_filesystem_class
+        cls = FilesystemJSONDecoder.try_resolve_fs_cls(dct)
+        if cls is None:
+            raise ValueError("Not a serialized AbstractFileSystem")
 
-        dic = json.loads(blob)
-        protocol = dic.pop("protocol")
-        try:
-            cls = _import_class(dic.pop("cls"))
-        except (ImportError, ValueError, RuntimeError, KeyError):
-            cls = get_filesystem_class(protocol)
-        return cls(*dic.pop("args", ()), **dic)
+        dct.pop("cls", None)
+        dct.pop("protocol", None)
+
+        return cls(*dct.pop("args", ()), **dct)
 
     def _get_pyarrow_filesystem(self):
         """
diff --git a/pyproject.toml b/pyproject.toml
index 93a0827..2dd6052 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,6 +29,7 @@ adl = ["adlfs"]
 arrow = ["pyarrow >= 1"]
 dask = ["dask", "distributed"]
 dev = ["ruff", "pre-commit"]
+doc = ["sphinx", "numpydoc", "sphinx-design", "sphinx-rtd-theme", "yarl"]
 dropbox = ["dropbox", "dropboxdrivefs", "requests"]
 entrypoints = []
 full = [
diff --git a/readthedocs.yml b/readthedocs.yml
index f44c84c..cd36aca 100644
--- a/readthedocs.yml
+++ b/readthedocs.yml
@@ -12,6 +12,8 @@ python:
   install:
     - method: pip
       path: .
+      extra_requirements:
+        - doc
 
 sphinx:
   configuration: docs/source/conf.py


Test patch:
diff --git a/fsspec/implementations/tests/test_cached.py b/fsspec/implementations/tests/test_cached.py
index f262867..091e6ca 100644
--- a/fsspec/implementations/tests/test_cached.py
+++ b/fsspec/implementations/tests/test_cached.py
@@ -873,7 +873,7 @@ def test_filecache_with_checks():
 
 @pytest.mark.parametrize("impl", ["filecache", "simplecache", "blockcache"])
 @pytest.mark.parametrize("fs", ["local", "multi"], indirect=["fs"])
-def test_takes_fs_instance(impl, fs):
+def test_filecache_takes_fs_instance(impl, fs):
     origin = tempfile.mkdtemp()
     data = b"test data"
     f1 = os.path.join(origin, "afile")
@@ -885,6 +885,15 @@ def test_takes_fs_instance(impl, fs):
     assert fs2.cat(f1) == data
 
 
+@pytest.mark.parametrize("impl", ["filecache", "simplecache", "blockcache"])
+@pytest.mark.parametrize("fs", ["local", "multi"], indirect=["fs"])
+def test_filecache_serialization(impl, fs):
+    fs1 = fsspec.filesystem(impl, fs=fs)
+    json1 = fs1.to_json()
+
+    assert fs1 is fsspec.AbstractFileSystem.from_json(json1)
+
+
 def test_add_file_to_cache_after_save(local_filecache):
     (data, original_file, cache_location, fs) = local_filecache
 
diff --git a/fsspec/implementations/tests/test_reference.py b/fsspec/implementations/tests/test_reference.py
index 3fe5e57..762d831 100644
--- a/fsspec/implementations/tests/test_reference.py
+++ b/fsspec/implementations/tests/test_reference.py
@@ -14,6 +14,10 @@ from fsspec.tests.conftest import data, realfile, reset_files, server, win  # no
 
 
 def test_simple(server):  # noqa: F811
+    # The dictionary in refs may be dumped with a different separator
+    # depending on whether json or ujson is imported
+    from fsspec.implementations.reference import json as json_impl
+
     refs = {
         "a": b"data",
         "b": (realfile, 0, 5),
@@ -28,12 +32,16 @@ def test_simple(server):  # noqa: F811
     assert fs.cat("b") == data[:5]
     assert fs.cat("c") == data[1 : 1 + 5]
     assert fs.cat("d") == b"hello"
-    assert fs.cat("e") == b'{"key": "value"}'
+    assert fs.cat("e") == json_impl.dumps(refs["e"]).encode("utf-8")
     with fs.open("d", "rt") as f:
         assert f.read(2) == "he"
 
 
 def test_simple_ver1(server):  # noqa: F811
+    # The dictionary in refs may be dumped with a different separator
+    # depending on whether json or ujson is imported
+    from fsspec.implementations.reference import json as json_impl
+
     in_data = {
         "version": 1,
         "refs": {
@@ -51,7 +59,7 @@ def test_simple_ver1(server):  # noqa: F811
     assert fs.cat("b") == data[:5]
     assert fs.cat("c") == data[1 : 1 + 5]
     assert fs.cat("d") == b"hello"
-    assert fs.cat("e") == b'{"key": "value"}'
+    assert fs.cat("e") == json_impl.dumps(in_data["refs"]["e"]).encode("utf-8")
     with fs.open("d", "rt") as f:
         assert f.read(2) == "he"
 
@@ -221,6 +229,8 @@ jdata = """{
 
 def test_spec1_expand():
     pytest.importorskip("jinja2")
+    from fsspec.implementations.reference import json as json_impl
+
     in_data = {
         "version": 1,
         "templates": {"u": "server.domain/path", "f": "{{c}}"},
@@ -244,6 +254,7 @@ def test_spec1_expand():
             "key2": ["http://{{u}}", 10000, 100],
             "key3": ["http://{{f(c='text')}}", 10000, 100],
             "key4": ["http://target_url"],
+            "key5": {"key": "value"},
         },
     }
     fs = fsspec.filesystem(
@@ -255,6 +266,7 @@ def test_spec1_expand():
         "key2": ["http://server.domain/path", 10000, 100],
         "key3": ["http://text", 10000, 100],
         "key4": ["http://target_url"],
+        "key5": json_impl.dumps(in_data["refs"]["key5"]),
         "gen_key0": ["http://server.domain/path_0", 1000, 1000],
         "gen_key1": ["http://server.domain/path_1", 2000, 1000],
         "gen_key2": ["http://server.domain/path_2", 3000, 1000],
@@ -267,6 +279,8 @@ def test_spec1_expand():
 
 def test_spec1_expand_simple():
     pytest.importorskip("jinja2")
+    from fsspec.implementations.reference import json as json_impl
+
     in_data = {
         "version": 1,
         "templates": {"u": "server.domain/path"},
@@ -274,6 +288,7 @@ def test_spec1_expand_simple():
             "key0": "base64:ZGF0YQ==",
             "key2": ["http://{{u}}", 10000, 100],
             "key4": ["http://target_url"],
+            "key5": {"key": "value"},
         },
     }
     fs = fsspec.filesystem("reference", fo=in_data, target_protocol="http")
@@ -286,6 +301,7 @@ def test_spec1_expand_simple():
     )
     assert fs.references["key2"] == ["http://not.org/p", 10000, 100]
     assert fs.cat("key0") == b"data"
+    assert fs.cat("key5") == json_impl.dumps(in_data["refs"]["key5"]).encode("utf-8")
 
 
 def test_spec1_gen_variants():
diff --git a/fsspec/tests/test_spec.py b/fsspec/tests/test_spec.py
index de23d78..27ea0f4 100644
--- a/fsspec/tests/test_spec.py
+++ b/fsspec/tests/test_spec.py
@@ -858,6 +858,66 @@ def test_json():
     assert DummyTestFS.from_json(outb) is b
 
 
+def test_json_path_attr():
+    a = DummyTestFS(1)
+    b = DummyTestFS(2, bar=Path("baz"))
+
+    outa = a.to_json()
+    outb = b.to_json()
+
+    assert json.loads(outb)  # is valid JSON
+    assert a != b
+    assert "bar" in outb
+
+    assert DummyTestFS.from_json(outa) is a
+    assert DummyTestFS.from_json(outb) is b
+
+
+def test_json_fs_attr():
+    a = DummyTestFS(1)
+    b = DummyTestFS(2, bar=a)
+
+    outa = a.to_json()
+    outb = b.to_json()
+
+    assert json.loads(outb)  # is valid JSON
+    assert a != b
+    assert "bar" in outb
+
+    assert DummyTestFS.from_json(outa) is a
+    assert DummyTestFS.from_json(outb) is b
+
+
+def test_from_dict_valid():
+    fs = DummyTestFS.from_dict({"cls": "fsspec.tests.test_spec.DummyTestFS"})
+    assert isinstance(fs, DummyTestFS)
+
+    fs = DummyTestFS.from_dict({"cls": "fsspec.tests.test_spec.DummyTestFS", "bar": 1})
+    assert fs.storage_options["bar"] == 1
+
+    fs = DummyTestFS.from_dict({"cls": "fsspec.implementations.local.LocalFileSystem"})
+    assert isinstance(fs, LocalFileSystem)
+
+    fs = DummyTestFS.from_dict(
+        {
+            "cls": "fsspec.implementations.local.LocalFileSystem",
+            "protocol": "local",
+        }
+    )
+    assert isinstance(fs, LocalFileSystem)
+
+
+def test_from_dict_invalid():
+    with pytest.raises(ValueError, match="Not a serialized AbstractFileSystem"):
+        DummyTestFS.from_dict({})
+
+    with pytest.raises(ValueError, match="Not a serialized AbstractFileSystem"):
+        DummyTestFS.from_dict({"cls": "pathlib.Path"})
+
+    with pytest.raises(ValueError, match="Not a serialized AbstractFileSystem"):
+        DummyTestFS.from_dict({"protocol": "local"})  # cls must be present
+
+
 def test_ls_from_cache():
     fs = DummyTestFS()
     uncached_results = fs.ls("top_level/second_level/", refresh=True)