diff --git a/virtualizarr/manifests/manifest.py b/virtualizarr/manifests/manifest.py index 666f48548..cab1f50d2 100644 --- a/virtualizarr/manifests/manifest.py +++ b/virtualizarr/manifests/manifest.py @@ -1,11 +1,22 @@ import json import re -from collections.abc import ItemsView, Iterable, Iterator, KeysView, ValuesView +from collections.abc import ( + AsyncIterator, + ItemsView, + Iterable, + Iterator, + KeysView, + ValuesView, +) from pathlib import PosixPath from typing import Any, Callable, NewType, Tuple, TypedDict, cast from urllib.parse import urlparse, urlunparse import numpy as np +import xarray as xr +from fsspec.asyn import AsyncFileSystem +from zarr.abc.store import RangeByteRequest, Store +from zarr.core.buffer import Buffer, BufferPrototype from virtualizarr.types import ChunkKey @@ -54,6 +65,110 @@ def with_validation( return ChunkEntry(path=path, offset=offset, length=length) +class ManifestStore(Store): + supports_writes: bool = False + supports_deletes: bool = False + supports_partial_writes: bool = False + supports_listing: bool = True + + fs: AsyncFileSystem + vds: xr.Dataset + + def __init__(self, fs, vds, read_only=True): + super().__init__(read_only=read_only) + self.fs = fs + self.vds = vds + + async def clear(self) -> None: + self.fs = None + self.vds = None + + def __str__(self) -> str: + return f"manifest://{id(self.vds)}" + + def __repr__(self) -> str: + return f"ManifestStore('{self}')" + + def __eq__(self, other: object) -> bool: + return ( + isinstance(other, type(self)) + and self.fs == other.fs + and xr.testing.assert_identical(self.vds, other.vds) + ) + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: None = None, # could this optionally accept a RangeByteRequest? + ) -> Buffer | None: + if not self._is_open: + await self._open() + print("key: ", key) + print("key split: ", key.split("/")) + array_name, _, chunk_key = key.split("/") + # TODO: is this the best way? + url, offset, length = self.vds[array_name].data.manifest.dict()[chunk_key] + value = prototype.buffer.from_bytes( + await self.fs._cat_file( + url, + start=offset, + end=offset + length, + ) + ) + return value + + # TODO: need a get_v3_array_metadata method + # to handle key="zarr.json" and return the metadata for the array + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, RangeByteRequest | None]], + ) -> list[Buffer | None]: + key_ranges = list(key_ranges) + paths: list[str] = [] + starts: list[int] = [] + stops: list[int] = [] + for key, _ in key_ranges: + array_name, _, chunk_key = key.split("/") + url, offset, length = self.vds[array_name].data.manifest.dict()[chunk_key] + paths.append(url) + starts.append(offset) + stops.append(offset + length) + res = await self.fs._cat_ranges(paths, starts, stops, on_error="return") + return [prototype.buffer.from_bytes(r) for r in res] + + async def exists(self, key: str) -> bool: + array_name, _, chunk_key = key.split("/") + url, _, _ = self.vds[array_name].data.manifest.dict()[chunk_key] + return await self.fs._exists(url) + + async def list(self) -> AsyncIterator[str]: + for array_name in self.vds.data_vars: + for chunk_key in self.vds[array_name].data.manifest: + yield f"{array_name}/{chunk_key}" + + async def list_prefix(self, prefix: str) -> AsyncIterator[str]: + raise NotImplementedError + + async def list_dir(self, prefix: str) -> AsyncIterator[str]: + raise NotImplementedError + + async def delete(self, key: str) -> None: + raise NotImplementedError + + async def set( + self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None + ) -> None: + raise NotImplementedError + + async def set_partial_values( + self, key_start_values: Iterable[tuple[str, int, bytes]] + ) -> None: + raise NotImplementedError + + def validate_and_normalize_path_to_uri(path: str, fs_root: str | None = None) -> str: """ Makes all paths into fully-qualified absolute URIs, or raises.