| import io |
| from pathlib import Path |
| from typing import Optional |
|
|
| from beaker import Beaker, ChecksumFailedError, DatasetNotFound, DatasetReadError |
|
|
| from .scheme_client import SchemeClient |
|
|
|
|
| class BeakerClient(SchemeClient): |
| scheme = ("beaker",) |
| recoverable_errors = SchemeClient.recoverable_errors + (DatasetReadError, ChecksumFailedError) |
|
|
| def __init__(self, resource: str) -> None: |
| super().__init__(resource) |
| self.beaker = Beaker.from_env() |
| |
| path = Path(resource.split("://")[1]) |
| if len(path.parts) < 2: |
| raise ValueError( |
| f"Invalid beaker resource URL '{resource}'. " |
| "Resources should be in the form 'beaker://{user_name}/{dataset_name}/{path_to_file}' " |
| "or beaker://{dataset_id}/{path_to_file}." |
| ) |
|
|
| try: |
| user, dataset_name, *filepath_parts = path.parts |
| self.dataset = self.beaker.dataset.get(f"{user}/{dataset_name}") |
| except DatasetNotFound: |
| dataset_id, *filepath_parts = path.parts |
| self.dataset = self.beaker.dataset.get(dataset_id) |
|
|
| self.filepath = "/".join(filepath_parts) |
| self.file_info = self.beaker.dataset.file_info(self.dataset, self.filepath) |
|
|
| def get_etag(self) -> Optional[str]: |
| return None if self.file_info.digest is None else str(self.file_info.digest) |
|
|
| def get_size(self) -> Optional[int]: |
| return self.file_info.size |
|
|
| def get_resource(self, temp_file: io.BufferedWriter) -> None: |
| for chunk in self.beaker.dataset.stream_file(self.dataset, self.filepath, quiet=True): |
| if chunk: |
| temp_file.write(chunk) |
|
|
