Extend JSON serialization capabilities (#1612)

---------

Co-authored-by: Martin Durant <[email protected]>

Author: DarkLight1337

fsspec/implementations/cached.py

@@ -434,6 +434,7 @@ def __getattribute__(self, item):
             "__hash__",
             "__eq__",
             "to_json",
+            "to_dict",
             "cache_size",
             "pipe_file",
             "pipe",
@@ -506,15 +507,6 @@ def __hash__(self):
             ^ 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

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
 

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

fsspec/registry.py

@@ -257,27 +257,33 @@ def get_filesystem_class(protocol):
 """
 
 
-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):

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 @@
 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 @@ def read_block(self, fn, offset, length, delimiter=None):
                 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):
         """