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