feat/memory store registry

changes/3679.feature.md

@@ -0,0 +1,3 @@
+Adds a new in-memory storage backend called `ManagedMemoryStore`. Instances of `ManagedMemoryStore`
+function similarly to `MemoryStore`, but instances of `ManagedMemoryStore` can be constructed from
+a URL like `memory://store`.

docs/user-guide/arrays.md

@@ -14,15 +14,14 @@ np.random.seed(0)
 
 ```python exec="true" session="arrays" source="above" result="ansi"
 import zarr
-store = zarr.storage.MemoryStore()
-z = zarr.create_array(store=store, shape=(10000, 10000), chunks=(1000, 1000), dtype='int32')
+z = zarr.create_array(store="memory://arrays-demo", shape=(10000, 10000), chunks=(1000, 1000), dtype='int32')
 print(z)
 ```
 
 The code above creates a 2-dimensional array of 32-bit integers with 10000 rows
 and 10000 columns, divided into chunks where each chunk has 1000 rows and 1000
-columns (and so there will be 100 chunks in total). The data is written to a
-[`zarr.storage.MemoryStore`][] (e.g. an in-memory dict). See
+columns (and so there will be 100 chunks in total). The data is written to an
+in-memory store (see [`zarr.storage.MemoryStore`][] for more details). See
 [Persistent arrays](#persistent-arrays) for details on storing arrays in other stores,
 and see [Data types](data_types.md) for an in-depth look at the data types supported
 by Zarr.

docs/user-guide/attributes.md

@@ -3,33 +3,32 @@
 Zarr arrays and groups support custom key/value attributes, which can be useful for
 storing application-specific metadata. For example:
 
-```python exec="true" session="arrays" source="above" result="ansi"
+```python exec="true" session="attributes" source="above" result="ansi"
 import zarr
-store = zarr.storage.MemoryStore()
-root = zarr.create_group(store=store)
+root = zarr.create_group(store="memory://attributes-demo")
 root.attrs['foo'] = 'bar'
 z = root.create_array(name='zzz', shape=(10000, 10000), dtype='int32')
 z.attrs['baz'] = 42
 z.attrs['qux'] = [1, 4, 7, 12]
 print(sorted(root.attrs))
 ```
 
-```python exec="true" session="arrays" source="above" result="ansi"
+```python exec="true" session="attributes" source="above" result="ansi"
 print('foo' in root.attrs)
 ```
 
-```python exec="true" session="arrays" source="above" result="ansi"
+```python exec="true" session="attributes" source="above" result="ansi"
 print(root.attrs['foo'])
 ```
-```python exec="true" session="arrays" source="above" result="ansi"
+```python exec="true" session="attributes" source="above" result="ansi"
 print(sorted(z.attrs))
 ```
 
-```python exec="true" session="arrays" source="above" result="ansi"
+```python exec="true" session="attributes" source="above" result="ansi"
 print(z.attrs['baz'])
 ```
 
-```python exec="true" session="arrays" source="above" result="ansi"
+```python exec="true" session="attributes" source="above" result="ansi"
 print(z.attrs['qux'])
 ```
 

docs/user-guide/consolidated_metadata.md

@@ -27,8 +27,7 @@ import zarr
 import warnings
 
 warnings.filterwarnings("ignore", category=UserWarning)
-store = zarr.storage.MemoryStore()
-group = zarr.create_group(store=store)
+group = zarr.create_group(store="memory://consolidated-metadata-demo")
 print(group)
 array = group.create_array(shape=(1,), name='a', dtype='float64')
 print(array)
@@ -45,7 +44,7 @@ print(array)
 ```
 
 ```python exec="true" session="consolidated_metadata" source="above" result="ansi"
-result = zarr.consolidate_metadata(store)
+result = zarr.consolidate_metadata("memory://consolidated-metadata-demo")
 print(result)
 ```
 
@@ -56,7 +55,7 @@ that can be used.:
 from pprint import pprint
 import io
 
-consolidated = zarr.open_group(store=store)
+consolidated = zarr.open_group(store="memory://consolidated-metadata-demo")
 consolidated_metadata = consolidated.metadata.consolidated_metadata.metadata
 
 # Note: pprint can be users without capturing the output regularly
@@ -76,7 +75,7 @@ With nested groups, the consolidated metadata is available on the children, recu
 ```python exec="true" session="consolidated_metadata" source="above" result="ansi"
 child = group.create_group('child', attributes={'kind': 'child'})
 grandchild = child.create_group('child', attributes={'kind': 'grandchild'})
-consolidated = zarr.consolidate_metadata(store)
+consolidated = zarr.consolidate_metadata("memory://consolidated-metadata-demo")
 
 output = io.StringIO()
 pprint(consolidated['child'].metadata.consolidated_metadata, stream=output, width=60)

docs/user-guide/gpu.md

@@ -20,9 +20,8 @@ buffers used internally by Zarr via `enable_gpu()`.
 import zarr
 import cupy as cp
 zarr.config.enable_gpu()
-store = zarr.storage.MemoryStore()
 z = zarr.create_array(
-    store=store, shape=(100, 100), chunks=(10, 10), dtype="float32",
+    store="memory://gpu-demo", shape=(100, 100), chunks=(10, 10), dtype="float32",
 )
 type(z[:10, :10])
 # cupy.ndarray

docs/user-guide/groups.md

@@ -8,8 +8,7 @@ To create a group, use the [`zarr.group`][] function:
 
 ```python exec="true" session="groups" source="above" result="ansi"
 import zarr
-store = zarr.storage.MemoryStore()
-root = zarr.create_group(store=store)
+root = zarr.create_group(store="memory://groups-demo")
 print(root)
 ```
 
@@ -105,8 +104,7 @@ Diagnostic information about arrays and groups is available via the `info`
 property. E.g.:
 
 ```python exec="true" session="groups" source="above" result="ansi"
-store = zarr.storage.MemoryStore()
-root = zarr.group(store=store)
+root = zarr.group(store="memory://diagnostics-demo")
 foo = root.create_group('foo')
 bar = foo.create_array(name='bar', shape=1000000, chunks=100000, dtype='int64')
 bar[:] = 42

src/zarr/storage/__init__.py

@@ -8,7 +8,7 @@
 from zarr.storage._fsspec import FsspecStore
 from zarr.storage._local import LocalStore
 from zarr.storage._logging import LoggingStore
-from zarr.storage._memory import GpuMemoryStore, MemoryStore
+from zarr.storage._memory import GpuMemoryStore, ManagedMemoryStore, MemoryStore
 from zarr.storage._obstore import ObjectStore
 from zarr.storage._wrapper import WrapperStore
 from zarr.storage._zip import ZipStore
@@ -18,6 +18,7 @@
     "GpuMemoryStore",
     "LocalStore",
     "LoggingStore",
+    "ManagedMemoryStore",
     "MemoryStore",
     "ObjectStore",
     "StoreLike",

src/zarr/storage/_common.py

@@ -17,8 +17,8 @@
 )
 from zarr.errors import ContainsArrayAndGroupError, ContainsArrayError, ContainsGroupError
 from zarr.storage._local import LocalStore
-from zarr.storage._memory import MemoryStore
-from zarr.storage._utils import normalize_path
+from zarr.storage._memory import ManagedMemoryStore, MemoryStore
+from zarr.storage._utils import _dereference_path, normalize_path
 
 _has_fsspec = importlib.util.find_spec("fsspec")
 if _has_fsspec:
@@ -30,18 +30,6 @@
     from zarr.core.buffer import BufferPrototype
 
 
-def _dereference_path(root: str, path: str) -> str:
-    if not isinstance(root, str):
-        msg = f"{root=} is not a string ({type(root)=})"  # type: ignore[unreachable]
-        raise TypeError(msg)
-    if not isinstance(path, str):
-        msg = f"{path=} is not a string ({type(path)=})"  # type: ignore[unreachable]
-        raise TypeError(msg)
-    root = root.rstrip("/")
-    path = f"{root}/{path}" if root else path
-    return path.rstrip("/")
-
-
 class StorePath:
     """
     Path-like interface for a Store.
@@ -341,8 +329,17 @@ async def make_store(
         return await LocalStore.open(root=store_like, mode=mode, read_only=_read_only)
 
     elif isinstance(store_like, str):
+        # Check for memory:// URLs first
+        if store_like.startswith("memory://"):
+            # Parse the URL to extract name and path
+            url_without_scheme = store_like[len("memory://") :]
+            parts = url_without_scheme.split("/", 1)
+            name = parts[0] if parts[0] else None
+            path = parts[1] if len(parts) > 1 else ""
+            # Create or get the store - ManagedMemoryStore handles both cases
+            return ManagedMemoryStore(name=name, path=path, read_only=_read_only)
         # Either an FSSpec URI or a local filesystem path
-        if _is_fsspec_uri(store_like):
+        elif _is_fsspec_uri(store_like):
             return FsspecStore.from_url(
                 store_like, storage_options=storage_options, read_only=_read_only
             )
@@ -418,6 +415,18 @@ async def make_store_path(
             "'path' was provided but is not used for FSMap store_like objects. Specify the path when creating the FSMap instance instead."
         )
 
+    elif isinstance(store_like, str) and store_like.startswith("memory://"):
+        # Handle memory:// URLs specially
+        # Parse the URL to extract name and path
+        _read_only = mode == "r"
+        url_without_scheme = store_like[len("memory://") :]
+        parts = url_without_scheme.split("/", 1)
+        name = parts[0] if parts[0] else None
+        url_path = parts[1] if len(parts) > 1 else ""
+        # Create or get the store - ManagedMemoryStore handles both cases
+        memory_store = ManagedMemoryStore(name=name, path=url_path, read_only=_read_only)
+        return await StorePath.open(memory_store, path=path_normalized, mode=mode)
+
     else:
         store = await make_store(store_like, mode=mode, storage_options=storage_options)
         return await StorePath.open(store, path=path_normalized, mode=mode)

src/zarr/storage/_fsspec.py

@@ -16,7 +16,7 @@
 )
 from zarr.core.buffer import Buffer
 from zarr.errors import ZarrUserWarning
-from zarr.storage._common import _dereference_path
+from zarr.storage._common import _dereference_path  # type: ignore[attr-defined]
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterable