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