from __future__ import annotations

import os

import threading

import weakref

from logging import getLogger

from typing import TYPE_CHECKING, Any, Self

from zarr.abc.store import ByteRequest, Store

from zarr.core.buffer import Buffer, gpu

from zarr.core.buffer.core import default_buffer_prototype

from zarr.core.common import concurrent_map

from zarr.storage._utils import (

_join_paths,

_normalize_byte_range_index,

normalize_path,

parse_store_url,

)

if TYPE_CHECKING:

from collections.abc import AsyncIterator, Iterable, MutableMapping

from zarr.core.buffer import BufferPrototype

logger = getLogger(__name__)

def _copy_buffer(value: Buffer) -> Buffer:

"""Copy `value` so the store does not retain the caller's memory.

Encoding a chunk can hand the store a zero-copy view of the user's array

(an uncompressed write is the common case), and unlike stores that

serialize on write, this one keeps whatever it is given alive in a dict.

Without this copy a later mutation of the user's array would rewrite

chunks already committed to the store.

"""

return type(value).from_array_like(value.as_array_like().copy())

class MemoryStore(Store):

"""

Store for local memory.

Parameters

----------

store_dict : dict

Initial data

read_only : bool

Whether the store is read-only

Attributes

----------

supports_writes

supports_deletes

supports_listing

Notes

-----

Writes copy the buffer they are given, so the store never aliases the

caller's memory. Buffers passed via `store_dict` are the caller's

responsibility and are stored as-is.

"""

supports_writes: bool = True

supports_deletes: bool = True

supports_listing: bool = True

_store_dict: MutableMapping[str, Buffer]

def __init__(

self,

store_dict: MutableMapping[str, Buffer] | None = None,

*,

read_only: bool = False,

) -> None:

super().__init__(read_only=read_only)

if store_dict is None:

store_dict = {}

self._store_dict = store_dict

def with_read_only(self, read_only: bool = False) -> MemoryStore:

# docstring inherited

return type(self)(

store_dict=self._store_dict,

read_only=read_only,

)

async def clear(self) -> None:

# docstring inherited

self._store_dict.clear()

def __str__(self) -> str:

return f"memory://{id(self._store_dict)}"

def __repr__(self) -> str:

return f"MemoryStore('{self}')"

def __eq__(self, other: object) -> bool:

return (

isinstance(other, type(self))

and self._store_dict == other._store_dict

and self.read_only == other.read_only

)

# -------------------------------------------------------------------

# Synchronous store methods

# -------------------------------------------------------------------

def get_sync(

self,

key: str,

*,

prototype: BufferPrototype | None = None,

byte_range: ByteRequest | None = None,

) -> Buffer | None:

if prototype is None:

prototype = default_buffer_prototype()

if not self._is_open:

self._is_open = True

try:

value = self._store_dict[key]

start, stop = _normalize_byte_range_index(value, byte_range)

return prototype.buffer.from_buffer(value[start:stop])

except KeyError:

return None

def set_sync(self, key: str, value: Buffer) -> None:

self._check_writable()

if not self._is_open:

self._is_open = True

if not isinstance(value, Buffer):

raise TypeError(

f"MemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."

)

self._store_dict[key] = _copy_buffer(value)

def delete_sync(self, key: str) -> None:

self._check_writable()

if not self._is_open:

self._is_open = True

try:

del self._store_dict[key]

except KeyError:

logger.debug("Key %s does not exist.", key)

async def get(

self,

key: str,

prototype: BufferPrototype | None = None,

byte_range: ByteRequest | None = None,

) -> Buffer | None:

# docstring inherited

if prototype is None:

prototype = default_buffer_prototype()

if not self._is_open:

await self._open()

try:

value = self._store_dict[key]

start, stop = _normalize_byte_range_index(value, byte_range)

return prototype.buffer.from_buffer(value[start:stop])

except KeyError:

return None

async def get_partial_values(

self,

prototype: BufferPrototype,

key_ranges: Iterable[tuple[str, ByteRequest | None]],

) -> list[Buffer | None]:

# docstring inherited

# All the key-ranges arguments goes with the same prototype

async def _get(key: str, byte_range: ByteRequest | None) -> Buffer | None:

return await self.get(key, prototype=prototype, byte_range=byte_range)

return await concurrent_map(key_ranges, _get, limit=None)

async def exists(self, key: str) -> bool:

# docstring inherited

return key in self._store_dict

async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None:

# docstring inherited

self._check_writable()

await self._ensure_open()

if not isinstance(value, Buffer):

raise TypeError(

f"MemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."

)

if byte_range is not None:

buf = self._store_dict[key]

buf[byte_range[0] : byte_range[1]] = value

self._store_dict[key] = buf

else:

self._store_dict[key] = _copy_buffer(value)

async def set_if_not_exists(self, key: str, value: Buffer) -> None:

# docstring inherited

self._check_writable()

await self._ensure_open()

self._store_dict.setdefault(key, _copy_buffer(value))

async def delete(self, key: str) -> None:

# docstring inherited

self._check_writable()

try:

del self._store_dict[key]

except KeyError:

logger.debug("Key %s does not exist.", key)

async def list(self) -> AsyncIterator[str]:

# docstring inherited

for key in self._store_dict:

yield key

async def list_prefix(self, prefix: str) -> AsyncIterator[str]:

# docstring inherited

# note: we materialize all dict keys into a list here so we can mutate the dict in-place (e.g. in delete_prefix)

for key in list(self._store_dict):

if key.startswith(prefix):

yield key

async def list_dir(self, prefix: str) -> AsyncIterator[str]:

# docstring inherited

prefix = prefix.rstrip("/")

if prefix == "":

keys_unique = {k.split("/")[0] for k in self._store_dict}

else:

# Our dictionary doesn't contain directory markers, but we want to include

# a pseudo directory when there's a nested item and we're listing an

# intermediate level.

keys_unique = {

key.removeprefix(f"{prefix}/").split("/")[0]

for key in self._store_dict

if key.startswith(f"{prefix}/") and key not in {prefix, f"{prefix}/"}

}

for key in keys_unique:

yield key

class GpuMemoryStore(MemoryStore):

"""

Store for GPU memory.

Stores every chunk in GPU memory irrespective of the original location.

The dictionary of buffers to initialize this memory store with *must* be

GPU Buffers.

Writing data to this store through ``.set`` will move the buffer to the GPU

if necessary.

Parameters

----------

store_dict : MutableMapping, optional

A mutable mapping with string keys and [zarr.core.buffer.gpu.Buffer][]

values.

read_only : bool

Whether to open the store in read-only mode.

"""

_store_dict: MutableMapping[str, gpu.Buffer] # type: ignore[assignment]

def __init__(

self,

store_dict: MutableMapping[str, gpu.Buffer] | None = None,

*,

read_only: bool = False,

) -> None:

super().__init__(store_dict=store_dict, read_only=read_only) # type: ignore[arg-type]

def __str__(self) -> str:

return f"gpumemory://{id(self._store_dict)}"

def __repr__(self) -> str:

return f"GpuMemoryStore('{self}')"

@classmethod

def from_dict(cls, store_dict: MutableMapping[str, Buffer]) -> Self:

"""

Create a GpuMemoryStore from a dictionary of buffers at any location.

The dictionary backing the newly created ``GpuMemoryStore`` will not be

the same as ``store_dict``.

Parameters

----------

store_dict : mapping

A mapping of strings keys to arbitrary Buffers. The buffer data

will be moved into a [`gpu.Buffer`][zarr.core.buffer.gpu.Buffer].

Returns

-------

GpuMemoryStore

"""

gpu_store_dict = {k: gpu.Buffer.from_buffer(v) for k, v in store_dict.items()}

return cls(gpu_store_dict)

async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None:

# docstring inherited

self._check_writable()

if not isinstance(value, Buffer):

raise TypeError(

f"GpuMemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."

)

# Convert to gpu.Buffer

gpu_value = value if isinstance(value, gpu.Buffer) else gpu.Buffer.from_buffer(value)

await super().set(key, gpu_value, byte_range=byte_range)

def set_sync(self, key: str, value: Buffer) -> None:

# docstring inherited

self._check_writable()

if not isinstance(value, Buffer):

raise TypeError(

f"GpuMemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."

)

# Convert to gpu.Buffer, mirroring `set` above: every value in this store's

# backing dict must be a gpu.Buffer, regardless of which API wrote it.

gpu_value = value if isinstance(value, gpu.Buffer) else gpu.Buffer.from_buffer(value)

super().set_sync(key, gpu_value)

# -----------------------------------------------------------------------------

# ManagedMemoryStore and its registry

# -----------------------------------------------------------------------------

# ManagedMemoryStore owns the lifecycle of its backing dict, enabling proper

# weakref-based tracking. This allows memory:// URLs to be resolved back to

# the store's dict within the same process.

class _ManagedStoreDict(dict[str, Buffer]):

"""

A dict subclass that supports weak references.

Regular dicts don't support weakrefs, but we need to track managed store dicts

in a WeakValueDictionary so they can be garbage collected when no longer

referenced. This subclass adds the necessary __weakref__ slot.

"""

__slots__ = ("__weakref__",)

class _ManagedStoreDictRegistry:

"""

Registry for managed store dicts.

This registry is the source of truth for managed store dicts. It creates

new dicts, tracks them via weak references, and looks them up by name.

"""

def __init__(self) -> None:

self._registry: weakref.WeakValueDictionary[str, _ManagedStoreDict] = (

weakref.WeakValueDictionary()

)

self._counter = 0

self._lock = threading.Lock()

def _generate_name(self) -> str:

"""Generate a unique name for a store.

Must be called while holding `self._lock`.

"""

name = str(self._counter)

self._counter += 1

return name

def get_or_create(self, name: str | None = None) -> tuple[_ManagedStoreDict, str]:

"""

Get an existing managed dict by name, or create a new one.

Thread-safe: uses a lock to prevent TOCTOU races between

checking for an existing entry and inserting a new one.

Parameters

----------

name : str | None

The name for the store. If None, a unique name is auto-generated.

If a store with this name already exists, returns the existing store.

Names cannot contain '/' characters.

Returns

-------

tuple[_ManagedStoreDict, str]

The store dict and its name.

Raises

------

ValueError

If the name contains '/' characters.

"""

with self._lock:

if name is None:

name = self._generate_name()

elif "/" in name:

raise ValueError(

f"Store name cannot contain '/': {name!r}. "

"Use the 'path' parameter to specify a path within the store."

)

existing = self._registry.get(name)

if existing is not None:

return existing, name

store_dict = _ManagedStoreDict()

self._registry[name] = store_dict

return store_dict, name

def get(self, name: str) -> _ManagedStoreDict | None:

"""

Look up a managed store dict by name.

Parameters

----------

name : str

The name of the store.

Returns

-------

_ManagedStoreDict | None

The store dict if found, None otherwise.

"""

return self._registry.get(name)

_managed_store_dict_registry = _ManagedStoreDictRegistry()

class ManagedMemoryStore(MemoryStore):

"""

A memory store that owns and manages the lifecycle of its backing dict.

Unlike ``MemoryStore`` which accepts any ``MutableMapping``, this store

creates and owns its backing dict internally. This enables proper lifecycle

management and allows the store to be looked up by its ``memory://`` URL

within the same process.

Parameters

----------

name : str | None

The name for this store, used in the ``memory://`` URL. If None, a unique

name is auto-generated. If a store with this name already exists, the

new store will share the same backing dict.

path : str

The root path for this store. All keys will be prefixed with this path.

read_only : bool

Whether the store is read-only.

Attributes

----------

name : str

The name of this store.

path : str

The root path of this store.

Notes

-----

The backing dict is tracked via weak references and will be garbage collected

when no ``ManagedMemoryStore`` instances reference it. URLs pointing to a

garbage-collected store will fail to resolve.

See Also

--------

MemoryStore : A memory store that accepts any MutableMapping.

Examples

--------

>>> store = ManagedMemoryStore(name="my-data")

>>> str(store)

'memory://my-data'

>>> # Later, resolve the URL back to the store's dict

>>> store2 = ManagedMemoryStore.from_url("memory://my-data")

>>> store2._store_dict is store._store_dict

True

>>> # Create a store with a path prefix

>>> store3 = ManagedMemoryStore.from_url("memory://my-data/subdir")

>>> store3.path

'subdir'

"""

_store_dict: _ManagedStoreDict

_name: str

path: str

def __init__(self, name: str | None = None, *, path: str = "", read_only: bool = False) -> None:

# Skip MemoryStore.__init__ and call Store.__init__ directly

# because we manage _store_dict via the registry, not via a user-supplied

# MutableMapping. If MemoryStore.__init__ ever adds logic beyond setting

# _store_dict, that logic must be replicated here.

Store.__init__(self, read_only=read_only)

# Get or create a managed dict from the registry

self._store_dict, self._name = _managed_store_dict_registry.get_or_create(name)

self.path = normalize_path(path)

def __str__(self) -> str:

return _join_paths([f"memory://{self._name}", self.path])

def __repr__(self) -> str:

return f"ManagedMemoryStore('{self}')"

def __eq__(self, other: object) -> bool:

return (

isinstance(other, type(self))

and self._store_dict is other._store_dict

and self.path == other.path

and self.read_only == other.read_only

)

@property

def name(self) -> str:

"""The name of this store, used in the memory:// URL."""

return self._name

@classmethod

def _from_managed_dict(

cls,

managed_dict: _ManagedStoreDict,

name: str,

*,

path: str = "",

read_only: bool = False,

) -> ManagedMemoryStore:

"""Internal: create a store from an existing managed dict."""

store = object.__new__(cls)

Store.__init__(store, read_only=read_only)

store._store_dict = managed_dict

store._name = name

store.path = normalize_path(path)

return store

def with_read_only(self, read_only: bool = False) -> ManagedMemoryStore:

# docstring inherited

return type(self)._from_managed_dict(

self._store_dict, self._name, path=self.path, read_only=read_only

)

@classmethod

def from_url(cls, url: str, *, read_only: bool = False) -> ManagedMemoryStore:

"""

Create a ManagedMemoryStore from a memory:// URL.

This looks up the backing dict in the process-wide registry and creates

a new store instance that shares the same dict.

Parameters

----------

url : str

A URL like "memory://my-store" or "memory://my-store/path/to/data"

identifying the store and optional path prefix.

read_only : bool

Whether the new store should be read-only.

Returns

-------

ManagedMemoryStore

A store sharing the same backing dict as the original.

Raises

------

ValueError

If the URL is not a valid memory:// URL or the store has been

garbage collected.

"""

parsed = parse_store_url(url)

if parsed.scheme != "memory":

raise ValueError(

f"Expected a 'memory://' URL, got scheme {parsed.scheme!r} in '{url}'."

)

name = parsed.name or ""

managed_dict = _managed_store_dict_registry.get(name)

if managed_dict is None:

raise ValueError(

f"Memory store not found for URL '{url}'. "

"The store may have been garbage collected."

)

return cls._from_managed_dict(managed_dict, name, path=parsed.path, read_only=read_only)

# Override MemoryStore methods to use path prefix and check process

def get_sync(

self,

key: str,

*,

prototype: BufferPrototype | None = None,

byte_range: ByteRequest | None = None,

) -> Buffer | None:

# docstring inherited

return super().get_sync(

_join_paths([self.path, key]), prototype=prototype, byte_range=byte_range

)

def set_sync(self, key: str, value: Buffer) -> None:

# docstring inherited

super().set_sync(_join_paths([self.path, key]), value)

def delete_sync(self, key: str) -> None:

# docstring inherited

super().delete_sync(_join_paths([self.path, key]))

async def get(

self,

key: str,

prototype: BufferPrototype | None = None,

byte_range: ByteRequest | None = None,

) -> Buffer | None:

# docstring inherited

return await super().get(

_join_paths([self.path, key]), prototype=prototype, byte_range=byte_range

)

# get_partial_values is intentionally NOT overridden here: MemoryStore.get_partial_values

# dispatches per-key through `self.get`, which already resolves to the override above.

# Re-prefixing the keys here as well would apply `self.path` twice.

async def exists(self, key: str) -> bool:

# docstring inherited

return await super().exists(_join_paths([self.path, key]))

async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None:

# docstring inherited

return await super().set(_join_paths([self.path, key]), value, byte_range=byte_range)

async def set_if_not_exists(self, key: str, value: Buffer) -> None:

# docstring inherited

return await super().set_if_not_exists(_join_paths([self.path, key]), value)

async def delete(self, key: str) -> None:

# docstring inherited

return await super().delete(_join_paths([self.path, key]))

async def list(self) -> AsyncIterator[str]:

# docstring inherited

prefix = f"{self.path}/" if self.path else ""

async for key in super().list():

if key.startswith(prefix):

yield key.removeprefix(prefix)

async def list_prefix(self, prefix: str) -> AsyncIterator[str]:

# docstring inherited

# Manual concatenation instead of _join_paths because we need "path/"

# as the prefix when prefix is empty (to list all keys under self.path)

full_prefix = f"{self.path}/{prefix}" if self.path else prefix

path_prefix = f"{self.path}/" if self.path else ""

async for key in super().list_prefix(full_prefix):

yield key.removeprefix(path_prefix)

async def list_dir(self, prefix: str) -> AsyncIterator[str]:

# docstring inherited

full_prefix = _join_paths([self.path, prefix])

async for key in super().list_dir(full_prefix):

yield key

def __reduce__(

self,

) -> tuple[type[ManagedMemoryStore], tuple[str | None], dict[str, Any]]:

"""

Support pickling of ManagedMemoryStore.

On unpickle, the store will reconnect to an existing store with the same

name if one exists in the registry, or create a new empty store otherwise.

Note that the backing dict data is NOT serialized - only the store's

identity (name, path, read_only) is preserved. If the original store has

been garbage collected, the unpickled store will have an empty dict.

The current process ID is preserved so that cross-process unpickling can be

detected and will raise an error at unpickle time.

"""

return (

self.__class__,

(self._name,),

{

"path": self.path,

"read_only": self.read_only,

"created_pid": os.getpid(),

},

)

def __setstate__(self, state: dict[str, Any]) -> None:

"""Restore state after unpickling.

The pickle protocol calls ``cls(name)`` (via ``__reduce__``'s args)

then ``__setstate__(state)``. ``__init__`` already set up

``_store_dict`` and ``_name`` from the registry — we just restore

path and read_only here.

"""

# Check for cross-process usage first, before mutating state

created_pid = state.get("created_pid")

if created_pid is not None and created_pid != os.getpid():

raise RuntimeError(

f"ManagedMemoryStore '{self._name}' was created in process {created_pid} "

f"but is being unpickled in process {os.getpid()}. "

"ManagedMemoryStore instances cannot be shared across processes because "

"their backing dict is not serialized. Use a persistent store (e.g., "

"LocalStore, ZipStore) for cross-process data sharing."

)

self.path = normalize_path(state.get("path", ""))

# Use the Store-level _read_only attribute directly because

# Store.__init__ was already called by __init__ during unpickling

self._read_only = state.get("read_only", False)