A Zarr store which wraps another, and simply redirects requests for zarr.json objects to another object name

alt_meta_store.py

from typing import TYPE_CHECKING
from zarr.storage import WrapperStore
from zarr.core.common import ZARR_JSON

if TYPE_CHECKING:
    from zarr.abc.store import (
        Store,
        RangeByteRequest,
        OffsetByteRequest,
        SuffixByteRequest,
        ByteRequest,
    )

    from zarr.core.buffer import Buffer, BufferPrototype
    from collections.abc import AsyncGenerator, Iterable


class AltMetaWrapperStore(WrapperStore):
    """A Zarr store which wraps another, and simply redirects requests for zarr.json objects to another object name"""
    def __init__(self, inner: Store, metadata_name: str) -> None:
        self.metadata_name = metadata_name
        super().__init__(WrapperStore(inner))

    def _map_key(self, key: str) -> str:
        """
        If the key is a zarr metadata key, replace the last path component with this store's `metadata_name`.
        """
        if key.endswith(ZARR_JSON) and (
            len(key) == len(ZARR_JSON) or key[-len(ZARR_JSON) - 1] == "/"
        ):
            key = key[: -len(ZARR_JSON)] + self.metadata_name
        return key

    async def get(
        self,
        key: str,
        prototype: BufferPrototype,
        byte_range: RangeByteRequest
        | OffsetByteRequest
        | SuffixByteRequest
        | None = None,
    ) -> Buffer | None:
        return await self._store.get(self._map_key(key), prototype, byte_range)

    async def get_partial_values(
        self,
        prototype: BufferPrototype,
        key_ranges: Iterable[
            tuple[str, RangeByteRequest | OffsetByteRequest | SuffixByteRequest | None]
        ],
    ) -> list[Buffer | None]:
        return await self._store.get_partial_values(
            prototype,
            [(self._map_key(key), byte_range) for key, byte_range in key_ranges],
        )

    async def exists(self, key: str) -> bool:
        return await self._store.exists(self._map_key(key))

    async def set(self, key: str, value: Buffer) -> None:
        return await self._store.set(self._map_key(key), value)

    async def delete(self, key: str) -> None:
        return await self._store.delete(self._map_key(key))

    async def set_if_not_exists(self, key: str, value: Buffer) -> None:
        return await self._store.set_if_not_exists(self._map_key(key), value)

    async def _get_many(
        self, requests: Iterable[tuple[str, BufferPrototype, ByteRequest | None]]
    ) -> AsyncGenerator[tuple[str, Buffer | None], None]:
        async for req in self._store._get_many(
            (self._map_key(k), b, r) for k, b, r in requests
        ):
            yield req