|
|
|
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 |
|
|
