| import io |
| import os |
| from collections.abc import Iterable |
| from typing import Optional, Union |
|
|
| from ..utils.file_utils import ( |
| SINGLE_FILE_COMPRESSION_PROTOCOLS, |
| ArchiveIterable, |
| FilesIterable, |
| _get_extraction_protocol, |
| _get_path_extension, |
| _prepare_path_and_storage_options, |
| is_relative_path, |
| url_or_path_join, |
| xbasename, |
| xdirname, |
| xet_parse, |
| xexists, |
| xgetsize, |
| xglob, |
| xgzip_open, |
| xisdir, |
| xisfile, |
| xjoin, |
| xlistdir, |
| xnumpy_load, |
| xopen, |
| xpandas_read_csv, |
| xpandas_read_excel, |
| xPath, |
| xpyarrow_parquet_read_table, |
| xrelpath, |
| xsio_loadmat, |
| xsplit, |
| xsplitext, |
| xwalk, |
| xxml_dom_minidom_parse, |
| ) |
| from ..utils.logging import get_logger |
| from ..utils.py_utils import map_nested |
| from .download_config import DownloadConfig |
|
|
|
|
| logger = get_logger(__name__) |
|
|
|
|
| class StreamingDownloadManager: |
| """ |
| Download manager that uses the "::" separator to navigate through (possibly remote) compressed archives. |
| Contrary to the regular `DownloadManager`, the `download` and `extract` methods don't actually download nor extract |
| data, but they rather return the path or url that could be opened using the `xopen` function which extends the |
| built-in `open` function to stream data from remote files. |
| """ |
|
|
| is_streaming = True |
|
|
| def __init__( |
| self, |
| dataset_name: Optional[str] = None, |
| data_dir: Optional[str] = None, |
| download_config: Optional[DownloadConfig] = None, |
| base_path: Optional[str] = None, |
| ): |
| self._dataset_name = dataset_name |
| self._data_dir = data_dir |
| self._base_path = base_path or os.path.abspath(".") |
| self.download_config = download_config or DownloadConfig() |
| self.downloaded_size = None |
| self.record_checksums = False |
|
|
| @property |
| def manual_dir(self): |
| return self._data_dir |
|
|
| def download(self, url_or_urls): |
| """Normalize URL(s) of files to stream data from. |
| This is the lazy version of `DownloadManager.download` for streaming. |
| |
| Args: |
| url_or_urls (`str` or `list` or `dict`): |
| URL(s) of files to stream data from. Each url is a `str`. |
| |
| Returns: |
| url(s): (`str` or `list` or `dict`), URL(s) to stream data from matching the given input url_or_urls. |
| |
| Example: |
| |
| ```py |
| >>> downloaded_files = dl_manager.download('https://storage.googleapis.com/seldon-datasets/sentence_polarity_v1/rt-polaritydata.tar.gz') |
| ``` |
| """ |
| url_or_urls = map_nested(self._download_single, url_or_urls, map_tuple=True) |
| return url_or_urls |
|
|
| def _download_single(self, urlpath: str) -> str: |
| urlpath = str(urlpath) |
| if is_relative_path(urlpath): |
| |
| urlpath = url_or_path_join(self._base_path, urlpath) |
| return urlpath |
|
|
| def extract(self, url_or_urls): |
| """Add extraction protocol for given url(s) for streaming. |
| |
| This is the lazy version of `DownloadManager.extract` for streaming. |
| |
| Args: |
| url_or_urls (`str` or `list` or `dict`): |
| URL(s) of files to stream data from. Each url is a `str`. |
| |
| Returns: |
| url(s): (`str` or `list` or `dict`), URL(s) to stream data from matching the given input `url_or_urls`. |
| |
| Example: |
| |
| ```py |
| >>> downloaded_files = dl_manager.download('https://storage.googleapis.com/seldon-datasets/sentence_polarity_v1/rt-polaritydata.tar.gz') |
| >>> extracted_files = dl_manager.extract(downloaded_files) |
| ``` |
| """ |
| urlpaths = map_nested(self._extract, url_or_urls, map_tuple=True) |
| return urlpaths |
|
|
| def _extract(self, urlpath: str) -> str: |
| urlpath = str(urlpath) |
| protocol = _get_extraction_protocol(urlpath, download_config=self.download_config) |
| |
| path = urlpath.split("::")[0] |
| extension = _get_path_extension(path) |
| if extension in ["tgz", "tar"] or path.endswith((".tar.gz", ".tar.bz2", ".tar.xz")): |
| raise NotImplementedError( |
| f"Extraction protocol for TAR archives like '{urlpath}' is not implemented in streaming mode. " |
| f"Please use `dl_manager.iter_archive` instead.\n\n" |
| f"Example usage:\n\n" |
| f"\turl = dl_manager.download(url)\n" |
| f"\ttar_archive_iterator = dl_manager.iter_archive(url)\n\n" |
| f"\tfor filename, file in tar_archive_iterator:\n" |
| f"\t\t..." |
| ) |
| if protocol is None: |
| |
| return urlpath |
| elif protocol in SINGLE_FILE_COMPRESSION_PROTOCOLS: |
| |
| inner_file = os.path.basename(urlpath.split("::")[0]) |
| inner_file = inner_file[: inner_file.rindex(".")] if "." in inner_file else inner_file |
| return f"{protocol}://{inner_file}::{urlpath}" |
| else: |
| return f"{protocol}://::{urlpath}" |
|
|
| def download_and_extract(self, url_or_urls): |
| """Prepare given `url_or_urls` for streaming (add extraction protocol). |
| |
| This is the lazy version of `DownloadManager.download_and_extract` for streaming. |
| |
| Is equivalent to: |
| |
| ``` |
| urls = dl_manager.extract(dl_manager.download(url_or_urls)) |
| ``` |
| |
| Args: |
| url_or_urls (`str` or `list` or `dict`): |
| URL(s) to stream from data from. Each url is a `str`. |
| |
| Returns: |
| url(s): (`str` or `list` or `dict`), URL(s) to stream data from matching the given input `url_or_urls`. |
| """ |
| return self.extract(self.download(url_or_urls)) |
|
|
| def iter_archive(self, urlpath_or_buf: Union[str, io.BufferedReader]) -> Iterable[tuple]: |
| """Iterate over files within an archive. |
| |
| Args: |
| urlpath_or_buf (`str` or `io.BufferedReader`): |
| Archive path or archive binary file object. |
| |
| Yields: |
| `tuple[str, io.BufferedReader]`: |
| 2-tuple (path_within_archive, file_object). |
| File object is opened in binary mode. |
| |
| Example: |
| |
| ```py |
| >>> archive = dl_manager.download('https://storage.googleapis.com/seldon-datasets/sentence_polarity_v1/rt-polaritydata.tar.gz') |
| >>> files = dl_manager.iter_archive(archive) |
| ``` |
| """ |
|
|
| if hasattr(urlpath_or_buf, "read"): |
| return ArchiveIterable.from_buf(urlpath_or_buf) |
| else: |
| return ArchiveIterable.from_urlpath(urlpath_or_buf, download_config=self.download_config) |
|
|
| def iter_files(self, urlpaths: Union[str, list[str]]) -> Iterable[str]: |
| """Iterate over files. |
| |
| Args: |
| urlpaths (`str` or `list` of `str`): |
| Root paths. |
| |
| Yields: |
| str: File URL path. |
| |
| Example: |
| |
| ```py |
| >>> files = dl_manager.download_and_extract('https://huggingface.co/datasets/beans/resolve/main/data/train.zip') |
| >>> files = dl_manager.iter_files(files) |
| ``` |
| """ |
| return FilesIterable.from_urlpaths(urlpaths, download_config=self.download_config) |
|
|
| def manage_extracted_files(self): |
| pass |
|
|
| def get_recorded_sizes_checksums(self): |
| pass |
|
|
