File size: 1,770 Bytes
6d63e5b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | 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()
# Beaker resources should be in the form "{user}/{dataset_name}/{path}/{to}/{file}"
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)
|