Spaces:
Running
on
T4
Running
on
T4
| # ========================================================== | |
| # 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') | |