|
|
|
|
|
|
|
|
|
import json |
|
import pickle |
|
from abc import ABCMeta, abstractmethod |
|
from pathlib import Path |
|
|
|
import yaml |
|
|
|
try: |
|
from yaml import CLoader as Loader, CDumper as Dumper |
|
except ImportError: |
|
from yaml import Loader, Dumper |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class BaseFileHandler(metaclass=ABCMeta): |
|
@abstractmethod |
|
def load_from_fileobj(self, file, **kwargs): |
|
pass |
|
|
|
@abstractmethod |
|
def dump_to_fileobj(self, obj, file, **kwargs): |
|
pass |
|
|
|
@abstractmethod |
|
def dump_to_str(self, obj, **kwargs): |
|
pass |
|
|
|
def load_from_path(self, filepath, mode="r", **kwargs): |
|
with open(filepath, mode) as f: |
|
return self.load_from_fileobj(f, **kwargs) |
|
|
|
def dump_to_path(self, obj, filepath, mode="w", **kwargs): |
|
with open(filepath, mode) as f: |
|
self.dump_to_fileobj(obj, f, **kwargs) |
|
|
|
|
|
class JsonHandler(BaseFileHandler): |
|
def load_from_fileobj(self, file): |
|
return json.load(file) |
|
|
|
def dump_to_fileobj(self, obj, file, **kwargs): |
|
json.dump(obj, file, **kwargs) |
|
|
|
def dump_to_str(self, obj, **kwargs): |
|
return json.dumps(obj, **kwargs) |
|
|
|
|
|
class PickleHandler(BaseFileHandler): |
|
def load_from_fileobj(self, file, **kwargs): |
|
return pickle.load(file, **kwargs) |
|
|
|
def load_from_path(self, filepath, **kwargs): |
|
return super(PickleHandler, self).load_from_path(filepath, mode="rb", **kwargs) |
|
|
|
def dump_to_str(self, obj, **kwargs): |
|
kwargs.setdefault("protocol", 2) |
|
return pickle.dumps(obj, **kwargs) |
|
|
|
def dump_to_fileobj(self, obj, file, **kwargs): |
|
kwargs.setdefault("protocol", 2) |
|
pickle.dump(obj, file, **kwargs) |
|
|
|
def dump_to_path(self, obj, filepath, **kwargs): |
|
super(PickleHandler, self).dump_to_path(obj, filepath, mode="wb", **kwargs) |
|
|
|
|
|
class YamlHandler(BaseFileHandler): |
|
def load_from_fileobj(self, file, **kwargs): |
|
kwargs.setdefault("Loader", Loader) |
|
return yaml.load(file, **kwargs) |
|
|
|
def dump_to_fileobj(self, obj, file, **kwargs): |
|
kwargs.setdefault("Dumper", Dumper) |
|
yaml.dump(obj, file, **kwargs) |
|
|
|
def dump_to_str(self, obj, **kwargs): |
|
kwargs.setdefault("Dumper", Dumper) |
|
return yaml.dump(obj, **kwargs) |
|
|
|
|
|
file_handlers = { |
|
"json": JsonHandler(), |
|
"yaml": YamlHandler(), |
|
"yml": YamlHandler(), |
|
"pickle": PickleHandler(), |
|
"pkl": PickleHandler(), |
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
def is_str(x): |
|
"""Whether the input is an string instance. |
|
|
|
Note: This method is deprecated since python 2 is no longer supported. |
|
""" |
|
return isinstance(x, str) |
|
|
|
|
|
def slload(file, file_format=None, **kwargs): |
|
"""Load data from json/yaml/pickle files. |
|
|
|
This method provides a unified api for loading data from serialized files. |
|
|
|
Args: |
|
file (str or :obj:`Path` or file-like object): Filename or a file-like |
|
object. |
|
file_format (str, optional): If not specified, the file format will be |
|
inferred from the file extension, otherwise use the specified one. |
|
Currently supported formats include "json", "yaml/yml" and |
|
"pickle/pkl". |
|
|
|
Returns: |
|
The content from the file. |
|
""" |
|
if isinstance(file, Path): |
|
file = str(file) |
|
if file_format is None and is_str(file): |
|
file_format = file.split(".")[-1] |
|
if file_format not in file_handlers: |
|
raise TypeError(f"Unsupported format: {file_format}") |
|
|
|
handler = file_handlers[file_format] |
|
if is_str(file): |
|
obj = handler.load_from_path(file, **kwargs) |
|
elif hasattr(file, "read"): |
|
obj = handler.load_from_fileobj(file, **kwargs) |
|
else: |
|
raise TypeError('"file" must be a filepath str or a file-object') |
|
return obj |
|
|
|
|
|
def sldump(obj, file=None, file_format=None, **kwargs): |
|
"""Dump data to json/yaml/pickle strings or files. |
|
|
|
This method provides a unified api for dumping data as strings or to files, |
|
and also supports custom arguments for each file format. |
|
|
|
Args: |
|
obj (any): The python object to be dumped. |
|
file (str or :obj:`Path` or file-like object, optional): If not |
|
specified, then the object is dump to a str, otherwise to a file |
|
specified by the filename or file-like object. |
|
file_format (str, optional): Same as :func:`load`. |
|
|
|
Returns: |
|
bool: True for success, False otherwise. |
|
""" |
|
if isinstance(file, Path): |
|
file = str(file) |
|
if file_format is None: |
|
if is_str(file): |
|
file_format = file.split(".")[-1] |
|
elif file is None: |
|
raise ValueError("file_format must be specified since file is None") |
|
if file_format not in file_handlers: |
|
raise TypeError(f"Unsupported format: {file_format}") |
|
|
|
handler = file_handlers[file_format] |
|
if file is None: |
|
return handler.dump_to_str(obj, **kwargs) |
|
elif is_str(file): |
|
handler.dump_to_path(obj, file, **kwargs) |
|
elif hasattr(file, "write"): |
|
handler.dump_to_fileobj(obj, file, **kwargs) |
|
else: |
|
raise TypeError('"file" must be a filename str or a file-object') |
|
|
