| |
| |
| import re |
| from abc import ABCMeta, abstractmethod |
| from pathlib import Path |
| from typing import Optional, Union |
|
|
|
|
| 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. |
| """ |
|
|
| @property |
| def name(self) -> str: |
| return self.__class__.__name__ |
|
|
| @abstractmethod |
| def get(self, filepath: str) -> bytes: |
| pass |
|
|
|
|
| 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. |
| conf_path (str, optional): Config path of Petrel client. Default: None. |
| `New in version 1.7.1`. |
| |
| 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 = False, |
| conf_path: str = None): |
| try: |
| from petrel_client import client |
| except ImportError: |
| raise ImportError('Please install petrel_client to enable ' |
| 'PetrelBackend.') |
|
|
| self._client = client.Client(conf_path=conf_path, 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, 1) |
| 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]) -> bytes: |
| """Read data from a given ``filepath`` with 'rb' mode. |
| |
| Args: |
| filepath (str or Path): Path to read data. |
| |
| Returns: |
| bytes: The loaded bytes. |
| """ |
| filepath = self._map_path(filepath) |
| filepath = self._format_path(filepath) |
| value = self._client.Get(filepath) |
| return value |
|
|
|
|
| class HardDiskBackend(BaseStorageBackend): |
| """Raw hard disks storage backend.""" |
|
|
| 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 |
|
|
