Spaces:
Sleeping
Sleeping
# Copyright (c) OpenMMLab. All rights reserved. | |
from io import BytesIO, StringIO | |
from pathlib import Path | |
from ..utils import is_list_of, is_str | |
from .file_client import FileClient | |
from .handlers import BaseFileHandler, JsonHandler, PickleHandler, YamlHandler | |
file_handlers = { | |
'json': JsonHandler(), | |
'yaml': YamlHandler(), | |
'yml': YamlHandler(), | |
'pickle': PickleHandler(), | |
'pkl': PickleHandler() | |
} | |
def load(file, file_format=None, file_client_args=None, **kwargs): | |
"""Load data from json/yaml/pickle files. | |
This method provides a unified api for loading data from serialized files. | |
Note: | |
In v1.3.16 and later, ``load`` supports loading data from serialized | |
files those can be storaged in different backends. | |
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". | |
file_client_args (dict, optional): Arguments to instantiate a | |
FileClient. See :class:`mmcv.fileio.FileClient` for details. | |
Default: None. | |
Examples: | |
>>> load('/path/of/your/file') # file is storaged in disk | |
>>> load('https://path/of/your/file') # file is storaged in Internet | |
>>> load('s3://path/of/your/file') # file is storaged in petrel | |
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): | |
file_client = FileClient.infer_client(file_client_args, file) | |
if handler.str_like: | |
with StringIO(file_client.get_text(file)) as f: | |
obj = handler.load_from_fileobj(f, **kwargs) | |
else: | |
with BytesIO(file_client.get(file)) as f: | |
obj = handler.load_from_fileobj(f, **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 dump(obj, file=None, file_format=None, file_client_args=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. | |
Note: | |
In v1.3.16 and later, ``dump`` supports dumping data as strings or to | |
files which is saved to different backends. | |
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 dumped to a str, otherwise to a file | |
specified by the filename or file-like object. | |
file_format (str, optional): Same as :func:`load`. | |
file_client_args (dict, optional): Arguments to instantiate a | |
FileClient. See :class:`mmcv.fileio.FileClient` for details. | |
Default: None. | |
Examples: | |
>>> dump('hello world', '/path/of/your/file') # disk | |
>>> dump('hello world', 's3://path/of/your/file') # ceph or petrel | |
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): | |
file_client = FileClient.infer_client(file_client_args, file) | |
if handler.str_like: | |
with StringIO() as f: | |
handler.dump_to_fileobj(obj, f, **kwargs) | |
file_client.put_text(f.getvalue(), file) | |
else: | |
with BytesIO() as f: | |
handler.dump_to_fileobj(obj, f, **kwargs) | |
file_client.put(f.getvalue(), file) | |
elif hasattr(file, 'write'): | |
handler.dump_to_fileobj(obj, file, **kwargs) | |
else: | |
raise TypeError('"file" must be a filename str or a file-object') | |
def _register_handler(handler, file_formats): | |
"""Register a handler for some file extensions. | |
Args: | |
handler (:obj:`BaseFileHandler`): Handler to be registered. | |
file_formats (str or list[str]): File formats to be handled by this | |
handler. | |
""" | |
if not isinstance(handler, BaseFileHandler): | |
raise TypeError( | |
f'handler must be a child of BaseFileHandler, not {type(handler)}') | |
if isinstance(file_formats, str): | |
file_formats = [file_formats] | |
if not is_list_of(file_formats, str): | |
raise TypeError('file_formats must be a str or a list of str') | |
for ext in file_formats: | |
file_handlers[ext] = handler | |
def register_handler(file_formats, **kwargs): | |
def wrap(cls): | |
_register_handler(cls(**kwargs), file_formats) | |
return cls | |
return wrap | |