|
|
|
import inspect |
|
import os |
|
import os.path as osp |
|
import re |
|
import tempfile |
|
import warnings |
|
from abc import ABCMeta, abstractmethod |
|
from contextlib import contextmanager |
|
from pathlib import Path |
|
from typing import Iterable, Iterator, Optional, Tuple, Union |
|
from urllib.request import urlopen |
|
|
|
import annotator.uniformer.mmcv as mmcv |
|
from annotator.uniformer.mmcv.utils.misc import has_method |
|
from annotator.uniformer.mmcv.utils.path import is_filepath |
|
|
|
|
|
class BaseStorageBackend(metaclass=ABCMeta): |
|
"""Abstract class of storage backends. |
|
|
|
All backends need to implement two apis: ``get()`` and ``get_text()``. |
|
``get()`` reads the file as a byte stream and ``get_text()`` reads the file |
|
as texts. |
|
""" |
|
|
|
|
|
_allow_symlink = False |
|
|
|
@property |
|
def name(self): |
|
return self.__class__.__name__ |
|
|
|
@property |
|
def allow_symlink(self): |
|
return self._allow_symlink |
|
|
|
@abstractmethod |
|
def get(self, filepath): |
|
pass |
|
|
|
@abstractmethod |
|
def get_text(self, filepath): |
|
pass |
|
|
|
|
|
class CephBackend(BaseStorageBackend): |
|
"""Ceph storage backend (for internal use). |
|
|
|
Args: |
|
path_mapping (dict|None): path mapping dict from local path to Petrel |
|
path. When ``path_mapping={'src': 'dst'}``, ``src`` in ``filepath`` |
|
will be replaced by ``dst``. Default: None. |
|
|
|
.. warning:: |
|
:class:`mmcv.fileio.file_client.CephBackend` will be deprecated, |
|
please use :class:`mmcv.fileio.file_client.PetrelBackend` instead. |
|
""" |
|
|
|
def __init__(self, path_mapping=None): |
|
try: |
|
import ceph |
|
except ImportError: |
|
raise ImportError('Please install ceph to enable CephBackend.') |
|
|
|
warnings.warn( |
|
'CephBackend will be deprecated, please use PetrelBackend instead') |
|
self._client = ceph.S3Client() |
|
assert isinstance(path_mapping, dict) or path_mapping is None |
|
self.path_mapping = path_mapping |
|
|
|
def get(self, filepath): |
|
filepath = str(filepath) |
|
if self.path_mapping is not None: |
|
for k, v in self.path_mapping.items(): |
|
filepath = filepath.replace(k, v) |
|
value = self._client.Get(filepath) |
|
value_buf = memoryview(value) |
|
return value_buf |
|
|
|
def get_text(self, filepath, encoding=None): |
|
raise NotImplementedError |
|
|
|
|
|
class PetrelBackend(BaseStorageBackend): |
|
"""Petrel storage backend (for internal use). |
|
|
|
PetrelBackend supports reading and writing data to multiple clusters. |
|
If the file path contains the cluster name, PetrelBackend will read data |
|
from specified cluster or write data to it. Otherwise, PetrelBackend will |
|
access the default cluster. |
|
|
|
Args: |
|
path_mapping (dict, optional): Path mapping dict from local path to |
|
Petrel path. When ``path_mapping={'src': 'dst'}``, ``src`` in |
|
``filepath`` will be replaced by ``dst``. Default: None. |
|
enable_mc (bool, optional): Whether to enable memcached support. |
|
Default: True. |
|
|
|
Examples: |
|
>>> filepath1 = 's3://path/of/file' |
|
>>> filepath2 = 'cluster-name:s3://path/of/file' |
|
>>> client = PetrelBackend() |
|
>>> client.get(filepath1) # get data from default cluster |
|
>>> client.get(filepath2) # get data from 'cluster-name' cluster |
|
""" |
|
|
|
def __init__(self, |
|
path_mapping: Optional[dict] = None, |
|
enable_mc: bool = True): |
|
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) |
|
assert isinstance(path_mapping, dict) or path_mapping is None |
|
self.path_mapping = path_mapping |
|
|
|
def _map_path(self, filepath: Union[str, Path]) -> str: |
|
"""Map ``filepath`` to a string path whose prefix will be replaced by |
|
:attr:`self.path_mapping`. |
|
|
|
Args: |
|
filepath (str): Path to be mapped. |
|
""" |
|
filepath = str(filepath) |
|
if self.path_mapping is not None: |
|
for k, v in self.path_mapping.items(): |
|
filepath = filepath.replace(k, v) |
|
return filepath |
|
|
|
def _format_path(self, filepath: str) -> str: |
|
"""Convert a ``filepath`` to standard format of petrel oss. |
|
|
|
If the ``filepath`` is concatenated by ``os.path.join``, in a Windows |
|
environment, the ``filepath`` will be the format of |
|
's3://bucket_name\\image.jpg'. By invoking :meth:`_format_path`, the |
|
above ``filepath`` will be converted to 's3://bucket_name/image.jpg'. |
|
|
|
Args: |
|
filepath (str): Path to be formatted. |
|
""" |
|
return re.sub(r'\\+', '/', filepath) |
|
|
|
def get(self, filepath: Union[str, Path]) -> memoryview: |
|
"""Read data from a given ``filepath`` with 'rb' mode. |
|
|
|
Args: |
|
filepath (str or Path): Path to read data. |
|
|
|
Returns: |
|
memoryview: A memory view of expected bytes object to avoid |
|
copying. The memoryview object can be converted to bytes by |
|
``value_buf.tobytes()``. |
|
""" |
|
filepath = self._map_path(filepath) |
|
filepath = self._format_path(filepath) |
|
value = self._client.Get(filepath) |
|
value_buf = memoryview(value) |
|
return value_buf |
|
|
|
def get_text(self, |
|
filepath: Union[str, Path], |
|
encoding: str = 'utf-8') -> str: |
|
"""Read data from a given ``filepath`` with 'r' mode. |
|
|
|
Args: |
|
filepath (str or Path): Path to read data. |
|
encoding (str): The encoding format used to open the ``filepath``. |
|
Default: 'utf-8'. |
|
|
|
Returns: |
|
str: Expected text reading from ``filepath``. |
|
""" |
|
return str(self.get(filepath), encoding=encoding) |
|
|
|
def put(self, obj: bytes, filepath: Union[str, Path]) -> None: |
|
"""Save data to a given ``filepath``. |
|
|
|
Args: |
|
obj (bytes): Data to be saved. |
|
filepath (str or Path): Path to write data. |
|
""" |
|
filepath = self._map_path(filepath) |
|
filepath = self._format_path(filepath) |
|
self._client.put(filepath, obj) |
|
|
|
def put_text(self, |
|
obj: str, |
|
filepath: Union[str, Path], |
|
encoding: str = 'utf-8') -> None: |
|
"""Save data to a given ``filepath``. |
|
|
|
Args: |
|
obj (str): Data to be written. |
|
filepath (str or Path): Path to write data. |
|
encoding (str): The encoding format used to encode the ``obj``. |
|
Default: 'utf-8'. |
|
""" |
|
self.put(bytes(obj, encoding=encoding), filepath) |
|
|
|
def remove(self, filepath: Union[str, Path]) -> None: |
|
"""Remove a file. |
|
|
|
Args: |
|
filepath (str or Path): Path to be removed. |
|
""" |
|
if not has_method(self._client, 'delete'): |
|
raise NotImplementedError( |
|
('Current version of Petrel Python SDK has not supported ' |
|
'the `delete` method, please use a higher version or dev' |
|
' branch instead.')) |
|
|
|
filepath = self._map_path(filepath) |
|
filepath = self._format_path(filepath) |
|
self._client.delete(filepath) |
|
|
|
def exists(self, filepath: Union[str, Path]) -> bool: |
|
"""Check whether a file path exists. |
|
|
|
Args: |
|
filepath (str or Path): Path to be checked whether exists. |
|
|
|
Returns: |
|
bool: Return ``True`` if ``filepath`` exists, ``False`` otherwise. |
|
""" |
|
if not (has_method(self._client, 'contains') |
|
and has_method(self._client, 'isdir')): |
|
raise NotImplementedError( |
|
('Current version of Petrel Python SDK has not supported ' |
|
'the `contains` and `isdir` methods, please use a higher' |
|
'version or dev branch instead.')) |
|
|
|
filepath = self._map_path(filepath) |
|
filepath = self._format_path(filepath) |
|
return self._client.contains(filepath) or self._client.isdir(filepath) |
|
|
|
def isdir(self, filepath: Union[str, Path]) -> bool: |
|
"""Check whether a file path is a directory. |
|
|
|
Args: |
|
filepath (str or Path): Path to be checked whether it is a |
|
directory. |
|
|
|
Returns: |
|
bool: Return ``True`` if ``filepath`` points to a directory, |
|
``False`` otherwise. |
|
""" |
|
if not has_method(self._client, 'isdir'): |
|
raise NotImplementedError( |
|
('Current version of Petrel Python SDK has not supported ' |
|
'the `isdir` method, please use a higher version or dev' |
|
' branch instead.')) |
|
|
|
filepath = self._map_path(filepath) |
|
filepath = self._format_path(filepath) |
|
return self._client.isdir(filepath) |
|
|
|
def isfile(self, filepath: Union[str, Path]) -> bool: |
|
"""Check whether a file path is a file. |
|
|
|
Args: |
|
filepath (str or Path): Path to be checked whether it is a file. |
|
|
|
Returns: |
|
bool: Return ``True`` if ``filepath`` points to a file, ``False`` |
|
otherwise. |
|
""" |
|
if not has_method(self._client, 'contains'): |
|
raise NotImplementedError( |
|
('Current version of Petrel Python SDK has not supported ' |
|
'the `contains` method, please use a higher version or ' |
|
'dev branch instead.')) |
|
|
|
filepath = self._map_path(filepath) |
|
filepath = self._format_path(filepath) |
|
return self._client.contains(filepath) |
|
|
|
def join_path(self, filepath: Union[str, Path], |
|
*filepaths: Union[str, Path]) -> str: |
|
"""Concatenate all file paths. |
|
|
|
Args: |
|
filepath (str or Path): Path to be concatenated. |
|
|
|
Returns: |
|
str: The result after concatenation. |
|
""" |
|
filepath = self._format_path(self._map_path(filepath)) |
|
if filepath.endswith('/'): |
|
filepath = filepath[:-1] |
|
formatted_paths = [filepath] |
|
for path in filepaths: |
|
formatted_paths.append(self._format_path(self._map_path(path))) |
|
return '/'.join(formatted_paths) |
|
|
|
@contextmanager |
|
def get_local_path(self, filepath: Union[str, Path]) -> Iterable[str]: |
|
"""Download a file from ``filepath`` and return a temporary path. |
|
|
|
``get_local_path`` is decorated by :meth:`contxtlib.contextmanager`. It |
|
can be called with ``with`` statement, and when exists from the |
|
``with`` statement, the temporary path will be released. |
|
|
|
Args: |
|
filepath (str | Path): Download a file from ``filepath``. |
|
|
|
Examples: |
|
>>> client = PetrelBackend() |
|
>>> # After existing from the ``with`` clause, |
|
>>> # the path will be removed |
|
>>> with client.get_local_path('s3://path/of/your/file') as path: |
|
... # do something here |
|
|
|
Yields: |
|
Iterable[str]: Only yield one temporary path. |
|
""" |
|
filepath = self._map_path(filepath) |
|
filepath = self._format_path(filepath) |
|
assert self.isfile(filepath) |
|
try: |
|
f = tempfile.NamedTemporaryFile(delete=False) |
|
f.write(self.get(filepath)) |
|
f.close() |
|
yield f.name |
|
finally: |
|
os.remove(f.name) |
|
|
|
def list_dir_or_file(self, |
|
dir_path: Union[str, Path], |
|
list_dir: bool = True, |
|
list_file: bool = True, |
|
suffix: Optional[Union[str, Tuple[str]]] = None, |
|
recursive: bool = False) -> Iterator[str]: |
|
"""Scan a directory to find the interested directories or files in |
|
arbitrary order. |
|
|
|
Note: |
|
Petrel has no concept of directories but it simulates the directory |
|
hierarchy in the filesystem through public prefixes. In addition, |
|
if the returned path ends with '/', it means the path is a public |
|
prefix which is a logical directory. |
|
|
|
Note: |
|
:meth:`list_dir_or_file` returns the path relative to ``dir_path``. |
|
In addition, the returned path of directory will not contains the |
|
suffix '/' which is consistent with other backends. |
|
|
|
Args: |
|
dir_path (str | Path): Path of the directory. |
|
list_dir (bool): List the directories. Default: True. |
|
list_file (bool): List the path of files. Default: True. |
|
suffix (str or tuple[str], optional): File suffix |
|
that we are interested in. Default: None. |
|
recursive (bool): If set to True, recursively scan the |
|
directory. Default: False. |
|
|
|
Yields: |
|
Iterable[str]: A relative path to ``dir_path``. |
|
""" |
|
if not has_method(self._client, 'list'): |
|
raise NotImplementedError( |
|
('Current version of Petrel Python SDK has not supported ' |
|
'the `list` method, please use a higher version or dev' |
|
' branch instead.')) |
|
|
|
dir_path = self._map_path(dir_path) |
|
dir_path = self._format_path(dir_path) |
|
if list_dir and suffix is not None: |
|
raise TypeError( |
|
'`list_dir` should be False when `suffix` is not None') |
|
|
|
if (suffix is not None) and not isinstance(suffix, (str, tuple)): |
|
raise TypeError('`suffix` must be a string or tuple of strings') |
|
|
|
|
|
|
|
if not dir_path.endswith('/'): |
|
dir_path += '/' |
|
|
|
root = dir_path |
|
|
|
def _list_dir_or_file(dir_path, list_dir, list_file, suffix, |
|
recursive): |
|
for path in self._client.list(dir_path): |
|
|
|
|
|
|
|
if path.endswith('/'): |
|
next_dir_path = self.join_path(dir_path, path) |
|
if list_dir: |
|
|
|
|
|
rel_dir = next_dir_path[len(root):-1] |
|
yield rel_dir |
|
if recursive: |
|
yield from _list_dir_or_file(next_dir_path, list_dir, |
|
list_file, suffix, |
|
recursive) |
|
else: |
|
absolute_path = self.join_path(dir_path, path) |
|
rel_path = absolute_path[len(root):] |
|
if (suffix is None |
|
or rel_path.endswith(suffix)) and list_file: |
|
yield rel_path |
|
|
|
return _list_dir_or_file(dir_path, list_dir, list_file, suffix, |
|
recursive) |
|
|
|
|
|
class MemcachedBackend(BaseStorageBackend): |
|
"""Memcached storage backend. |
|
|
|
Attributes: |
|
server_list_cfg (str): Config file for memcached server list. |
|
client_cfg (str): Config file for memcached client. |
|
sys_path (str | None): Additional path to be appended to `sys.path`. |
|
Default: None. |
|
""" |
|
|
|
def __init__(self, server_list_cfg, client_cfg, sys_path=None): |
|
if sys_path is not None: |
|
import sys |
|
sys.path.append(sys_path) |
|
try: |
|
import mc |
|
except ImportError: |
|
raise ImportError( |
|
'Please install memcached to enable MemcachedBackend.') |
|
|
|
self.server_list_cfg = server_list_cfg |
|
self.client_cfg = client_cfg |
|
self._client = mc.MemcachedClient.GetInstance(self.server_list_cfg, |
|
self.client_cfg) |
|
|
|
self._mc_buffer = mc.pyvector() |
|
|
|
def get(self, filepath): |
|
filepath = str(filepath) |
|
import mc |
|
self._client.Get(filepath, self._mc_buffer) |
|
value_buf = mc.ConvertBuffer(self._mc_buffer) |
|
return value_buf |
|
|
|
def get_text(self, filepath, encoding=None): |
|
raise NotImplementedError |
|
|
|
|
|
class LmdbBackend(BaseStorageBackend): |
|
"""Lmdb storage backend. |
|
|
|
Args: |
|
db_path (str): Lmdb database path. |
|
readonly (bool, optional): Lmdb environment parameter. If True, |
|
disallow any write operations. Default: True. |
|
lock (bool, optional): Lmdb environment parameter. If False, when |
|
concurrent access occurs, do not lock the database. Default: False. |
|
readahead (bool, optional): Lmdb environment parameter. If False, |
|
disable the OS filesystem readahead mechanism, which may improve |
|
random read performance when a database is larger than RAM. |
|
Default: False. |
|
|
|
Attributes: |
|
db_path (str): Lmdb database path. |
|
""" |
|
|
|
def __init__(self, |
|
db_path, |
|
readonly=True, |
|
lock=False, |
|
readahead=False, |
|
**kwargs): |
|
try: |
|
import lmdb |
|
except ImportError: |
|
raise ImportError('Please install lmdb to enable LmdbBackend.') |
|
|
|
self.db_path = str(db_path) |
|
self._client = lmdb.open( |
|
self.db_path, |
|
readonly=readonly, |
|
lock=lock, |
|
readahead=readahead, |
|
**kwargs) |
|
|
|
def get(self, filepath): |
|
"""Get values according to the filepath. |
|
|
|
Args: |
|
filepath (str | obj:`Path`): Here, filepath is the lmdb key. |
|
""" |
|
filepath = str(filepath) |
|
with self._client.begin(write=False) as txn: |
|
value_buf = txn.get(filepath.encode('ascii')) |
|
return value_buf |
|
|
|
def get_text(self, filepath, encoding=None): |
|
raise NotImplementedError |
|
|
|
|
|
class HardDiskBackend(BaseStorageBackend): |
|
"""Raw hard disks storage backend.""" |
|
|
|
_allow_symlink = True |
|
|
|
def get(self, filepath: Union[str, Path]) -> bytes: |
|
"""Read data from a given ``filepath`` with 'rb' mode. |
|
|
|
Args: |
|
filepath (str or Path): Path to read data. |
|
|
|
Returns: |
|
bytes: Expected bytes object. |
|
""" |
|
with open(filepath, 'rb') as f: |
|
value_buf = f.read() |
|
return value_buf |
|
|
|
def get_text(self, |
|
filepath: Union[str, Path], |
|
encoding: str = 'utf-8') -> str: |
|
"""Read data from a given ``filepath`` with 'r' mode. |
|
|
|
Args: |
|
filepath (str or Path): Path to read data. |
|
encoding (str): The encoding format used to open the ``filepath``. |
|
Default: 'utf-8'. |
|
|
|
Returns: |
|
str: Expected text reading from ``filepath``. |
|
""" |
|
with open(filepath, 'r', encoding=encoding) as f: |
|
value_buf = f.read() |
|
return value_buf |
|
|
|
def put(self, obj: bytes, filepath: Union[str, Path]) -> None: |
|
"""Write data to a given ``filepath`` with 'wb' mode. |
|
|
|
Note: |
|
``put`` will create a directory if the directory of ``filepath`` |
|
does not exist. |
|
|
|
Args: |
|
obj (bytes): Data to be written. |
|
filepath (str or Path): Path to write data. |
|
""" |
|
mmcv.mkdir_or_exist(osp.dirname(filepath)) |
|
with open(filepath, 'wb') as f: |
|
f.write(obj) |
|
|
|
def put_text(self, |
|
obj: str, |
|
filepath: Union[str, Path], |
|
encoding: str = 'utf-8') -> None: |
|
"""Write data to a given ``filepath`` with 'w' mode. |
|
|
|
Note: |
|
``put_text`` will create a directory if the directory of |
|
``filepath`` does not exist. |
|
|
|
Args: |
|
obj (str): Data to be written. |
|
filepath (str or Path): Path to write data. |
|
encoding (str): The encoding format used to open the ``filepath``. |
|
Default: 'utf-8'. |
|
""" |
|
mmcv.mkdir_or_exist(osp.dirname(filepath)) |
|
with open(filepath, 'w', encoding=encoding) as f: |
|
f.write(obj) |
|
|
|
def remove(self, filepath: Union[str, Path]) -> None: |
|
"""Remove a file. |
|
|
|
Args: |
|
filepath (str or Path): Path to be removed. |
|
""" |
|
os.remove(filepath) |
|
|
|
def exists(self, filepath: Union[str, Path]) -> bool: |
|
"""Check whether a file path exists. |
|
|
|
Args: |
|
filepath (str or Path): Path to be checked whether exists. |
|
|
|
Returns: |
|
bool: Return ``True`` if ``filepath`` exists, ``False`` otherwise. |
|
""" |
|
return osp.exists(filepath) |
|
|
|
def isdir(self, filepath: Union[str, Path]) -> bool: |
|
"""Check whether a file path is a directory. |
|
|
|
Args: |
|
filepath (str or Path): Path to be checked whether it is a |
|
directory. |
|
|
|
Returns: |
|
bool: Return ``True`` if ``filepath`` points to a directory, |
|
``False`` otherwise. |
|
""" |
|
return osp.isdir(filepath) |
|
|
|
def isfile(self, filepath: Union[str, Path]) -> bool: |
|
"""Check whether a file path is a file. |
|
|
|
Args: |
|
filepath (str or Path): Path to be checked whether it is a file. |
|
|
|
Returns: |
|
bool: Return ``True`` if ``filepath`` points to a file, ``False`` |
|
otherwise. |
|
""" |
|
return osp.isfile(filepath) |
|
|
|
def join_path(self, filepath: Union[str, Path], |
|
*filepaths: Union[str, Path]) -> str: |
|
"""Concatenate all file paths. |
|
|
|
Join one or more filepath components intelligently. The return value |
|
is the concatenation of filepath and any members of *filepaths. |
|
|
|
Args: |
|
filepath (str or Path): Path to be concatenated. |
|
|
|
Returns: |
|
str: The result of concatenation. |
|
""" |
|
return osp.join(filepath, *filepaths) |
|
|
|
@contextmanager |
|
def get_local_path( |
|
self, filepath: Union[str, Path]) -> Iterable[Union[str, Path]]: |
|
"""Only for unified API and do nothing.""" |
|
yield filepath |
|
|
|
def list_dir_or_file(self, |
|
dir_path: Union[str, Path], |
|
list_dir: bool = True, |
|
list_file: bool = True, |
|
suffix: Optional[Union[str, Tuple[str]]] = None, |
|
recursive: bool = False) -> Iterator[str]: |
|
"""Scan a directory to find the interested directories or files in |
|
arbitrary order. |
|
|
|
Note: |
|
:meth:`list_dir_or_file` returns the path relative to ``dir_path``. |
|
|
|
Args: |
|
dir_path (str | Path): Path of the directory. |
|
list_dir (bool): List the directories. Default: True. |
|
list_file (bool): List the path of files. Default: True. |
|
suffix (str or tuple[str], optional): File suffix |
|
that we are interested in. Default: None. |
|
recursive (bool): If set to True, recursively scan the |
|
directory. Default: False. |
|
|
|
Yields: |
|
Iterable[str]: A relative path to ``dir_path``. |
|
""" |
|
if list_dir and suffix is not None: |
|
raise TypeError('`suffix` should be None when `list_dir` is True') |
|
|
|
if (suffix is not None) and not isinstance(suffix, (str, tuple)): |
|
raise TypeError('`suffix` must be a string or tuple of strings') |
|
|
|
root = dir_path |
|
|
|
def _list_dir_or_file(dir_path, list_dir, list_file, suffix, |
|
recursive): |
|
for entry in os.scandir(dir_path): |
|
if not entry.name.startswith('.') and entry.is_file(): |
|
rel_path = osp.relpath(entry.path, root) |
|
if (suffix is None |
|
or rel_path.endswith(suffix)) and list_file: |
|
yield rel_path |
|
elif osp.isdir(entry.path): |
|
if list_dir: |
|
rel_dir = osp.relpath(entry.path, root) |
|
yield rel_dir |
|
if recursive: |
|
yield from _list_dir_or_file(entry.path, list_dir, |
|
list_file, suffix, |
|
recursive) |
|
|
|
return _list_dir_or_file(dir_path, list_dir, list_file, suffix, |
|
recursive) |
|
|
|
|
|
class HTTPBackend(BaseStorageBackend): |
|
"""HTTP and HTTPS storage bachend.""" |
|
|
|
def get(self, filepath): |
|
value_buf = urlopen(filepath).read() |
|
return value_buf |
|
|
|
def get_text(self, filepath, encoding='utf-8'): |
|
value_buf = urlopen(filepath).read() |
|
return value_buf.decode(encoding) |
|
|
|
@contextmanager |
|
def get_local_path(self, filepath: str) -> Iterable[str]: |
|
"""Download a file from ``filepath``. |
|
|
|
``get_local_path`` is decorated by :meth:`contxtlib.contextmanager`. It |
|
can be called with ``with`` statement, and when exists from the |
|
``with`` statement, the temporary path will be released. |
|
|
|
Args: |
|
filepath (str): Download a file from ``filepath``. |
|
|
|
Examples: |
|
>>> client = HTTPBackend() |
|
>>> # After existing from the ``with`` clause, |
|
>>> # the path will be removed |
|
>>> with client.get_local_path('http://path/of/your/file') as path: |
|
... # do something here |
|
""" |
|
try: |
|
f = tempfile.NamedTemporaryFile(delete=False) |
|
f.write(self.get(filepath)) |
|
f.close() |
|
yield f.name |
|
finally: |
|
os.remove(f.name) |
|
|
|
|
|
class FileClient: |
|
"""A general file client to access files in different backends. |
|
|
|
The client loads a file or text in a specified backend from its path |
|
and returns it as a binary or text file. There are two ways to choose a |
|
backend, the name of backend and the prefix of path. Although both of them |
|
can be used to choose a storage backend, ``backend`` has a higher priority |
|
that is if they are all set, the storage backend will be chosen by the |
|
backend argument. If they are all `None`, the disk backend will be chosen. |
|
Note that It can also register other backend accessor with a given name, |
|
prefixes, and backend class. In addition, We use the singleton pattern to |
|
avoid repeated object creation. If the arguments are the same, the same |
|
object will be returned. |
|
|
|
Args: |
|
backend (str, optional): The storage backend type. Options are "disk", |
|
"ceph", "memcached", "lmdb", "http" and "petrel". Default: None. |
|
prefix (str, optional): The prefix of the registered storage backend. |
|
Options are "s3", "http", "https". Default: None. |
|
|
|
Examples: |
|
>>> # only set backend |
|
>>> file_client = FileClient(backend='petrel') |
|
>>> # only set prefix |
|
>>> file_client = FileClient(prefix='s3') |
|
>>> # set both backend and prefix but use backend to choose client |
|
>>> file_client = FileClient(backend='petrel', prefix='s3') |
|
>>> # if the arguments are the same, the same object is returned |
|
>>> file_client1 = FileClient(backend='petrel') |
|
>>> file_client1 is file_client |
|
True |
|
|
|
Attributes: |
|
client (:obj:`BaseStorageBackend`): The backend object. |
|
""" |
|
|
|
_backends = { |
|
'disk': HardDiskBackend, |
|
'ceph': CephBackend, |
|
'memcached': MemcachedBackend, |
|
'lmdb': LmdbBackend, |
|
'petrel': PetrelBackend, |
|
'http': HTTPBackend, |
|
} |
|
|
|
|
|
|
|
|
|
_overridden_backends = set() |
|
_prefix_to_backends = { |
|
's3': PetrelBackend, |
|
'http': HTTPBackend, |
|
'https': HTTPBackend, |
|
} |
|
_overridden_prefixes = set() |
|
|
|
_instances = {} |
|
|
|
def __new__(cls, backend=None, prefix=None, **kwargs): |
|
if backend is None and prefix is None: |
|
backend = 'disk' |
|
if backend is not None and backend not in cls._backends: |
|
raise ValueError( |
|
f'Backend {backend} is not supported. Currently supported ones' |
|
f' are {list(cls._backends.keys())}') |
|
if prefix is not None and prefix not in cls._prefix_to_backends: |
|
raise ValueError( |
|
f'prefix {prefix} is not supported. Currently supported ones ' |
|
f'are {list(cls._prefix_to_backends.keys())}') |
|
|
|
|
|
|
|
arg_key = f'{backend}:{prefix}' |
|
for key, value in kwargs.items(): |
|
arg_key += f':{key}:{value}' |
|
|
|
|
|
if (arg_key in cls._instances |
|
and backend not in cls._overridden_backends |
|
and prefix not in cls._overridden_prefixes): |
|
_instance = cls._instances[arg_key] |
|
else: |
|
|
|
_instance = super().__new__(cls) |
|
if backend is not None: |
|
_instance.client = cls._backends[backend](**kwargs) |
|
else: |
|
_instance.client = cls._prefix_to_backends[prefix](**kwargs) |
|
|
|
cls._instances[arg_key] = _instance |
|
|
|
return _instance |
|
|
|
@property |
|
def name(self): |
|
return self.client.name |
|
|
|
@property |
|
def allow_symlink(self): |
|
return self.client.allow_symlink |
|
|
|
@staticmethod |
|
def parse_uri_prefix(uri: Union[str, Path]) -> Optional[str]: |
|
"""Parse the prefix of a uri. |
|
|
|
Args: |
|
uri (str | Path): Uri to be parsed that contains the file prefix. |
|
|
|
Examples: |
|
>>> FileClient.parse_uri_prefix('s3://path/of/your/file') |
|
's3' |
|
|
|
Returns: |
|
str | None: Return the prefix of uri if the uri contains '://' |
|
else ``None``. |
|
""" |
|
assert is_filepath(uri) |
|
uri = str(uri) |
|
if '://' not in uri: |
|
return None |
|
else: |
|
prefix, _ = uri.split('://') |
|
|
|
|
|
if ':' in prefix: |
|
_, prefix = prefix.split(':') |
|
return prefix |
|
|
|
@classmethod |
|
def infer_client(cls, |
|
file_client_args: Optional[dict] = None, |
|
uri: Optional[Union[str, Path]] = None) -> 'FileClient': |
|
"""Infer a suitable file client based on the URI and arguments. |
|
|
|
Args: |
|
file_client_args (dict, optional): Arguments to instantiate a |
|
FileClient. Default: None. |
|
uri (str | Path, optional): Uri to be parsed that contains the file |
|
prefix. Default: None. |
|
|
|
Examples: |
|
>>> uri = 's3://path/of/your/file' |
|
>>> file_client = FileClient.infer_client(uri=uri) |
|
>>> file_client_args = {'backend': 'petrel'} |
|
>>> file_client = FileClient.infer_client(file_client_args) |
|
|
|
Returns: |
|
FileClient: Instantiated FileClient object. |
|
""" |
|
assert file_client_args is not None or uri is not None |
|
if file_client_args is None: |
|
file_prefix = cls.parse_uri_prefix(uri) |
|
return cls(prefix=file_prefix) |
|
else: |
|
return cls(**file_client_args) |
|
|
|
@classmethod |
|
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)}') |
|
if not inspect.isclass(backend): |
|
raise TypeError( |
|
f'backend should be a class but got {type(backend)}') |
|
if not issubclass(backend, BaseStorageBackend): |
|
raise TypeError( |
|
f'backend {backend} is not a subclass of BaseStorageBackend') |
|
if not force and name in cls._backends: |
|
raise KeyError( |
|
f'{name} is already registered as a storage backend, ' |
|
'add "force=True" if you want to override it') |
|
|
|
if name in cls._backends and force: |
|
cls._overridden_backends.add(name) |
|
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: |
|
cls._prefix_to_backends[prefix] = backend |
|
elif (prefix in cls._prefix_to_backends) and force: |
|
cls._overridden_prefixes.add(prefix) |
|
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, prefixes=None): |
|
"""Register a backend to FileClient. |
|
|
|
This method can be used as a normal class method or a decorator. |
|
|
|
.. code-block:: python |
|
|
|
class NewBackend(BaseStorageBackend): |
|
|
|
def get(self, filepath): |
|
return filepath |
|
|
|
def get_text(self, filepath): |
|
return filepath |
|
|
|
FileClient.register_backend('new', NewBackend) |
|
|
|
or |
|
|
|
.. code-block:: python |
|
|
|
@FileClient.register_backend('new') |
|
class NewBackend(BaseStorageBackend): |
|
|
|
def get(self, filepath): |
|
return filepath |
|
|
|
def get_text(self, filepath): |
|
return filepath |
|
|
|
Args: |
|
name (str): The name of the registered backend. |
|
backend (class, optional): The backend class to be registered, |
|
which must be a subclass of :class:`BaseStorageBackend`. |
|
When this method is used as a decorator, backend is None. |
|
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], optional): The prefixes |
|
of the registered storage backend. Default: None. |
|
`New in version 1.3.15.` |
|
""" |
|
if backend is not None: |
|
cls._register_backend( |
|
name, backend, force=force, prefixes=prefixes) |
|
return |
|
|
|
def _register(backend_cls): |
|
cls._register_backend( |
|
name, backend_cls, force=force, prefixes=prefixes) |
|
return backend_cls |
|
|
|
return _register |
|
|
|
def get(self, filepath: Union[str, Path]) -> Union[bytes, memoryview]: |
|
"""Read data from a given ``filepath`` with 'rb' mode. |
|
|
|
Note: |
|
There are two types of return values for ``get``, one is ``bytes`` |
|
and the other is ``memoryview``. The advantage of using memoryview |
|
is that you can avoid copying, and if you want to convert it to |
|
``bytes``, you can use ``.tobytes()``. |
|
|
|
Args: |
|
filepath (str or Path): Path to read data. |
|
|
|
Returns: |
|
bytes | memoryview: Expected bytes object or a memory view of the |
|
bytes object. |
|
""" |
|
return self.client.get(filepath) |
|
|
|
def get_text(self, filepath: Union[str, Path], encoding='utf-8') -> str: |
|
"""Read data from a given ``filepath`` with 'r' mode. |
|
|
|
Args: |
|
filepath (str or Path): Path to read data. |
|
encoding (str): The encoding format used to open the ``filepath``. |
|
Default: 'utf-8'. |
|
|
|
Returns: |
|
str: Expected text reading from ``filepath``. |
|
""" |
|
return self.client.get_text(filepath, encoding) |
|
|
|
def put(self, obj: bytes, filepath: Union[str, Path]) -> None: |
|
"""Write data to a given ``filepath`` with 'wb' mode. |
|
|
|
Note: |
|
``put`` should create a directory if the directory of ``filepath`` |
|
does not exist. |
|
|
|
Args: |
|
obj (bytes): Data to be written. |
|
filepath (str or Path): Path to write data. |
|
""" |
|
self.client.put(obj, filepath) |
|
|
|
def put_text(self, obj: str, filepath: Union[str, Path]) -> None: |
|
"""Write data to a given ``filepath`` with 'w' mode. |
|
|
|
Note: |
|
``put_text`` should create a directory if the directory of |
|
``filepath`` does not exist. |
|
|
|
Args: |
|
obj (str): Data to be written. |
|
filepath (str or Path): Path to write data. |
|
encoding (str, optional): The encoding format used to open the |
|
`filepath`. Default: 'utf-8'. |
|
""" |
|
self.client.put_text(obj, filepath) |
|
|
|
def remove(self, filepath: Union[str, Path]) -> None: |
|
"""Remove a file. |
|
|
|
Args: |
|
filepath (str, Path): Path to be removed. |
|
""" |
|
self.client.remove(filepath) |
|
|
|
def exists(self, filepath: Union[str, Path]) -> bool: |
|
"""Check whether a file path exists. |
|
|
|
Args: |
|
filepath (str or Path): Path to be checked whether exists. |
|
|
|
Returns: |
|
bool: Return ``True`` if ``filepath`` exists, ``False`` otherwise. |
|
""" |
|
return self.client.exists(filepath) |
|
|
|
def isdir(self, filepath: Union[str, Path]) -> bool: |
|
"""Check whether a file path is a directory. |
|
|
|
Args: |
|
filepath (str or Path): Path to be checked whether it is a |
|
directory. |
|
|
|
Returns: |
|
bool: Return ``True`` if ``filepath`` points to a directory, |
|
``False`` otherwise. |
|
""" |
|
return self.client.isdir(filepath) |
|
|
|
def isfile(self, filepath: Union[str, Path]) -> bool: |
|
"""Check whether a file path is a file. |
|
|
|
Args: |
|
filepath (str or Path): Path to be checked whether it is a file. |
|
|
|
Returns: |
|
bool: Return ``True`` if ``filepath`` points to a file, ``False`` |
|
otherwise. |
|
""" |
|
return self.client.isfile(filepath) |
|
|
|
def join_path(self, filepath: Union[str, Path], |
|
*filepaths: Union[str, Path]) -> str: |
|
"""Concatenate all file paths. |
|
|
|
Join one or more filepath components intelligently. The return value |
|
is the concatenation of filepath and any members of *filepaths. |
|
|
|
Args: |
|
filepath (str or Path): Path to be concatenated. |
|
|
|
Returns: |
|
str: The result of concatenation. |
|
""" |
|
return self.client.join_path(filepath, *filepaths) |
|
|
|
@contextmanager |
|
def get_local_path(self, filepath: Union[str, Path]) -> Iterable[str]: |
|
"""Download data from ``filepath`` and write the data to local path. |
|
|
|
``get_local_path`` is decorated by :meth:`contxtlib.contextmanager`. It |
|
can be called with ``with`` statement, and when exists from the |
|
``with`` statement, the temporary path will be released. |
|
|
|
Note: |
|
If the ``filepath`` is a local path, just return itself. |
|
|
|
.. warning:: |
|
``get_local_path`` is an experimental interface that may change in |
|
the future. |
|
|
|
Args: |
|
filepath (str or Path): Path to be read data. |
|
|
|
Examples: |
|
>>> file_client = FileClient(prefix='s3') |
|
>>> with file_client.get_local_path('s3://bucket/abc.jpg') as path: |
|
... # do something here |
|
|
|
Yields: |
|
Iterable[str]: Only yield one path. |
|
""" |
|
with self.client.get_local_path(str(filepath)) as local_path: |
|
yield local_path |
|
|
|
def list_dir_or_file(self, |
|
dir_path: Union[str, Path], |
|
list_dir: bool = True, |
|
list_file: bool = True, |
|
suffix: Optional[Union[str, Tuple[str]]] = None, |
|
recursive: bool = False) -> Iterator[str]: |
|
"""Scan a directory to find the interested directories or files in |
|
arbitrary order. |
|
|
|
Note: |
|
:meth:`list_dir_or_file` returns the path relative to ``dir_path``. |
|
|
|
Args: |
|
dir_path (str | Path): Path of the directory. |
|
list_dir (bool): List the directories. Default: True. |
|
list_file (bool): List the path of files. Default: True. |
|
suffix (str or tuple[str], optional): File suffix |
|
that we are interested in. Default: None. |
|
recursive (bool): If set to True, recursively scan the |
|
directory. Default: False. |
|
|
|
Yields: |
|
Iterable[str]: A relative path to ``dir_path``. |
|
""" |
|
yield from self.client.list_dir_or_file(dir_path, list_dir, list_file, |
|
suffix, recursive) |
|
|