[Feature] Loading objects from different backends and dumping objects to different backends

mmcv/fileio/file_client.py

@@ -1,6 +1,8 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import inspect
 from abc import ABCMeta, abstractmethod
+from pathlib import Path
+from typing import Optional, Union
 from urllib.request import urlopen
 
 
@@ -49,7 +51,7 @@ def get(self, filepath):
         value_buf = memoryview(value)
         return value_buf
 
-    def get_text(self, filepath):
+    def get_text(self, filepath, encoding=None):
         raise NotImplementedError
 
 
@@ -61,20 +63,26 @@ class PetrelBackend(BaseStorageBackend):
             path. When `path_mapping={'src': 'dst'}`, `src` in `filepath` will
             be replaced by `dst`. Default: None.
         enable_mc (bool): whether to enable memcached support. Default: True.
+        enable_multi_cluster (bool): Whether to enable multi clusters.
+            Default: False.
     """
 
-    def __init__(self, path_mapping=None, enable_mc=True):
+    def __init__(self,
+                 path_mapping: Optional[dict] = None,
+                 enable_mc: bool = True,
+                 enable_multi_cluster: bool = False):
         try:
             from petrel_client import client
         except ImportError:
             raise ImportError('Please install petrel_client to enable '
                               'PetrelBackend.')
 
-        self._client = client.Client(enable_mc=enable_mc)
+        self._client = client.Client(
+            enable_mc=enable_mc, enable_multi_cluster=enable_multi_cluster)
         assert isinstance(path_mapping, dict) or path_mapping is None
         self.path_mapping = path_mapping
 
-    def get(self, filepath):
+    def get(self, filepath: Union[str, Path]) -> memoryview:
         filepath = str(filepath)
         if self.path_mapping is not None:
             for k, v in self.path_mapping.items():
@@ -83,8 +91,23 @@ def get(self, filepath):
         value_buf = memoryview(value)
         return value_buf
 
-    def get_text(self, filepath):
-        raise NotImplementedError
+    def get_text(self,
+                 filepath: Union[str, Path],
+                 encoding: str = 'utf-8') -> str:
+        return str(self.get(filepath), encoding=encoding)
+
+    def put(self, obj: bytes, filepath: Union[str, Path]) -> None:
+        filepath = str(filepath)
+        if self.path_mapping is not None:
+            for k, v in self.path_mapping.items():
+                filepath = filepath.replace(k, v)
+        self._client.put(filepath, obj)
+
+    def put_text(self,
+                 obj: str,
+                 filepath: Union[str, Path],
+                 encoding: str = 'utf-8') -> None:
+        self.put(bytes(obj, encoding=encoding), filepath)
 
 
 class MemcachedBackend(BaseStorageBackend):
@@ -121,7 +144,7 @@ def get(self, filepath):
         value_buf = mc.ConvertBuffer(self._mc_buffer)
         return value_buf
 
-    def get_text(self, filepath):
+    def get_text(self, filepath, encoding=None):
         raise NotImplementedError
 
 
@@ -173,7 +196,7 @@ def get(self, filepath):
             value_buf = txn.get(filepath.encode('ascii'))
         return value_buf
 
-    def get_text(self, filepath):
+    def get_text(self, filepath, encoding=None):
         raise NotImplementedError
 
 
@@ -186,12 +209,22 @@ def get(self, filepath):
             value_buf = f.read()
         return value_buf
 
-    def get_text(self, filepath):
+    def get_text(self, filepath, encoding='utf-8'):
         filepath = str(filepath)
-        with open(filepath, 'r') as f:
+        with open(filepath, 'r', encoding=encoding) as f:
             value_buf = f.read()
         return value_buf
 
+    def put(self, obj, filepath):
+        filepath = str(filepath)
+        with open(filepath, 'wb') as f:
+            f.write(obj)
+
+    def put_text(self, obj, filepath, encoding='utf-8'):
+        filepath = str(filepath)
+        with open(filepath, 'w', encoding=encoding) as f:
+            f.write(obj)
+
 
 class HTTPBackend(BaseStorageBackend):
     """HTTP and HTTPS storage bachend."""
@@ -200,9 +233,9 @@ def get(self, filepath):
         value_buf = urlopen(filepath).read()
         return value_buf
 
-    def get_text(self, filepath):
+    def get_text(self, filepath, encoding='utf-8'):
         value_buf = urlopen(filepath).read()
-        return value_buf.decode('utf-8')
+        return value_buf.decode(encoding)
 
 
 class FileClient:
@@ -212,9 +245,20 @@ class FileClient:
     and return it as a binary file. it can also register other backend
     accessor with a given name and backend class.
 
-    Attributes:
+    Args:
         backend (str): The storage backend type. Options are "disk", "ceph",
-            "memcached", "lmdb" and "http".
+            "memcached", "lmdb", "http" and "petrel". Default: None.
+        prefixes (str or list[str] or tuple[str]): The prefixes of the
+            registered storage backend. Both backend and prefixes can be used
+            to choose a storage backend, but backend has a higher priority that
+            is if they are all set, the storage backend will be chosen by the
+            backend rather than prefixes. If backend and prefixes are all
+            `None`. The dist backend is be chosen. Default: None.
+
+    .. versionadd:: 1.3.14
+            The *prefixes* parameter.
+
+    Attributes:
         client (:obj:`BaseStorageBackend`): The backend object.
     """
 
@@ -226,17 +270,52 @@ class FileClient:
         'petrel': PetrelBackend,
         'http': HTTPBackend,
     }
+    _prefix_to_backends = {
+        's3': PetrelBackend,
+        'http': HTTPBackend,
+        'https': HTTPBackend,
+    }
 
-    def __init__(self, backend='disk', **kwargs):
-        if backend not in self._backends:
+    def __init__(self, backend=None, prefixes=None, **kwargs):
+        if backend is None and prefixes is None:
+            backend = 'disk'
+        if backend is not None and backend not in self._backends:
             raise ValueError(
                 f'Backend {backend} is not supported. Currently supported ones'
                 f' are {list(self._backends.keys())}')
-        self.backend = backend
-        self.client = self._backends[backend](**kwargs)
+        if prefixes is not None:
+            if isinstance(prefixes, str):
+                prefixes = [prefixes]
+            else:
+                assert isinstance(prefixes, (list, tuple))
+
+            if not set(prefixes).issubset(self._prefix_to_backends.keys()):
+                raise ValueError(
+                    f'prefixes {prefixes} is not supported. Currently '
+                    'supported ones are '
+                    f'{list(self._prefix_to_backends.keys())}')
+
+        if backend is not None:
+            self.client = self._backends[backend](**kwargs)
+        else:
+            for prefix in prefixes:
+                self.client = self._prefix_to_backends[prefix](**kwargs)
+                break
+
+    @staticmethod
+    def parse_uri_prefix(uri):
+        uri = str(uri)
+        if '://' not in uri:
+            return None
+        else:
+            prefix, _ = uri.split('://')
+            # clusterName:s3://
+            if ':' in prefix:
+                _, prefix = prefix.split(':')
+            return prefix
 
     @classmethod
-    def _register_backend(cls, name, backend, force=False):
+    def _register_backend(cls, name, backend, force=False, prefixes=None):
         if not isinstance(name, str):
             raise TypeError('the backend name should be a string, '
                             f'but got {type(name)}')
@@ -252,9 +331,21 @@ def _register_backend(cls, name, backend, force=False):
                 'add "force=True" if you want to override it')
 
         cls._backends[name] = backend
+        if prefixes is not None:
+            if isinstance(prefixes, str):
+                prefixes = [prefixes]
+            else:
+                assert isinstance(prefixes, (list, tuple))
+            for prefix in prefixes:
+                if (prefix not in cls._prefix_to_backends) or force:
+                    cls._prefix_to_backends[prefix] = backend
+                else:
+                    raise KeyError(
+                        f'{prefix} is already registered as a storage backend,'
+                        ' add "force=True" if you want to override it')
 
     @classmethod
-    def register_backend(cls, name, backend=None, force=False):
+    def register_backend(cls, name, backend=None, force=False, prefixes=None):
         """Register a backend to FileClient.
 
         This method can be used as a normal class method or a decorator.
@@ -292,19 +383,32 @@ def get_text(self, filepath):
                 Defaults to None.
             force (bool, optional): Whether to override the backend if the name
                 has already been registered. Defaults to False.
+            prefixes (str or list[str] or tuple[str]): The prefix of the
+                registered storage backend.
+
+        .. versionadd:: 1.3.14
+            The *prefixes* parameter.
         """
         if backend is not None:
-            cls._register_backend(name, backend, force=force)
+            cls._register_backend(
+                name, backend, force=force, prefixes=prefixes)
             return
 
         def _register(backend_cls):
-            cls._register_backend(name, backend_cls, force=force)
+            cls._register_backend(
+                name, backend_cls, force=force, prefixes=prefixes)
             return backend_cls
 
         return _register
 
     def get(self, filepath):
         return self.client.get(filepath)
 
-    def get_text(self, filepath):
-        return self.client.get_text(filepath)
+    def get_text(self, filepath, encoding='utf-8'):
+        return self.client.get_text(filepath, encoding)
+
+    def put(self, obj, filepath):
+        self.client.put(obj, filepath)
+
+    def put_text(self, obj, filepath):
+        self.client.put_text(obj, filepath)

mmcv/fileio/handlers/base.py

@@ -3,6 +3,7 @@
 
 
 class BaseFileHandler(metaclass=ABCMeta):
+    str_like_obj = True
 
     @abstractmethod
     def load_from_fileobj(self, file, **kwargs):

mmcv/fileio/handlers/pickle_handler.py

@@ -6,6 +6,8 @@
 
 class PickleHandler(BaseFileHandler):
 
+    str_like_obj = False
+
     def load_from_fileobj(self, file, **kwargs):
         return pickle.load(file, **kwargs)
 

mmcv/fileio/io.py

@@ -1,7 +1,9 @@
 # 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 = {
@@ -13,7 +15,7 @@
 }
 
 
-def load(file, file_format=None, **kwargs):
+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.
@@ -25,6 +27,8 @@ def load(file, file_format=None, **kwargs):
             inferred from the file extension, otherwise use the specified one.
             Currently supported formats include "json", "yaml/yml" and
             "pickle/pkl".
+        file_client_args (dict): Arguments to instantiate a FileClient.
+            See :class:`mmcv.fileio.FileClient` for details. Default: None.
 
     Returns:
         The content from the file.
@@ -36,17 +40,28 @@ def load(file, file_format=None, **kwargs):
     if file_format not in file_handlers:
         raise TypeError(f'Unsupported format: {file_format}')
 
+    if file_client_args is None:
+        file_prefix = FileClient.parse_uri_prefix(file)
+        client = FileClient(prefixes=file_prefix)
+    else:
+        client = FileClient(**file_client_args)
+
     handler = file_handlers[file_format]
     if is_str(file):
-        obj = handler.load_from_path(file, **kwargs)
+        if handler.str_like_obj:
+            with StringIO(client.get_text(file)) as f:
+                obj = handler.load_from_fileobj(f, **kwargs)
+        else:
+            with BytesIO(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, **kwargs):
+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,
@@ -58,7 +73,8 @@ def dump(obj, file=None, file_format=None, **kwargs):
             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`.
-
+        file_client_args (dict): Arguments to instantiate a FileClient.
+            See :class:`mmcv.fileio.FileClient` for details. Default: None.
     Returns:
         bool: True for success, False otherwise.
     """
@@ -73,11 +89,24 @@ def dump(obj, file=None, file_format=None, **kwargs):
     if file_format not in file_handlers:
         raise TypeError(f'Unsupported format: {file_format}')
 
+    if file_client_args is None:
+        file_prefix = FileClient.parse_uri_prefix(file)
+        client = FileClient(prefixes=file_prefix)
+    else:
+        client = FileClient(**file_client_args)
+
     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)
+        if handler.str_like_obj:
+            f = StringIO()
+            handler.dump_to_fileobj(obj, f, **kwargs)
+            client.put_text(f.getvalue(), file)
+        else:
+            f = BytesIO()
+            handler.dump_to_fileobj(obj, f, **kwargs)
+            client.put(f.getvalue(), file)
     elif hasattr(file, 'write'):
         handler.dump_to_fileobj(obj, file, **kwargs)
     else: