Spaces:
Running
Running
# ========================================================== | |
# Modified from mmcv | |
# ========================================================== | |
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 | |
# =========================== | |
# Rigister handler | |
# =========================== | |
class BaseFileHandler(metaclass=ABCMeta): | |
def load_from_fileobj(self, file, **kwargs): | |
pass | |
def dump_to_fileobj(self, obj, file, **kwargs): | |
pass | |
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(), | |
} | |
# =========================== | |
# load and dump | |
# =========================== | |
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') | |