You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
82 lines
2.6 KiB
82 lines
2.6 KiB
5 months ago
|
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
|