Spaces:
Running
on
Zero
Running
on
Zero
# Copyright (c) Meta Platforms, Inc. and affiliates. | |
# All rights reserved. | |
# | |
# This source code is licensed under the license found in the | |
# LICENSE file in the root directory of this source tree. | |
from mmengine.utils import is_list_of | |
from .base import BaseFileHandler | |
from .json_handler import JsonHandler | |
from .pickle_handler import PickleHandler | |
from .yaml_handler import YamlHandler | |
file_handlers = { | |
'json': JsonHandler(), | |
'yaml': YamlHandler(), | |
'yml': YamlHandler(), | |
'pickle': PickleHandler(), | |
'pkl': PickleHandler(), | |
} | |
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 | |