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, parse_store_url
 
 _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.
@@ -305,14 +293,18 @@ async def make_store(
     """
     from zarr.storage._fsspec import FsspecStore  # circular import
 
-    if (
-        not (isinstance(store_like, str) and _is_fsspec_uri(store_like))
-        and storage_options is not None
-    ):
-        raise TypeError(
-            "'storage_options' was provided but unused. "
-            "'storage_options' is only used when the store is passed as an FSSpec URI string.",
-        )
+    # Check if storage_options is valid for this store_like
+    if storage_options is not None:
+        is_fsspec_uri = False
+        if isinstance(store_like, str):
+            parsed = parse_store_url(store_like)
+            # fsspec URIs have a scheme that's not memory, file, or empty
+            is_fsspec_uri = parsed.scheme not in ("", "memory", "file")
+        if not is_fsspec_uri:
+            raise TypeError(
+                "'storage_options' was provided but unused. "
+                "'storage_options' is only used when the store is passed as an FSSpec URI string.",
+            )
 
     assert mode in (None, "r", "r+", "a", "w", "w-")
     _read_only = mode == "r"
@@ -341,14 +333,19 @@ async def make_store(
         return await LocalStore.open(root=store_like, mode=mode, read_only=_read_only)
 
     elif isinstance(store_like, str):
-        # Either an FSSpec URI or a local filesystem path
-        if _is_fsspec_uri(store_like):
+        parsed = parse_store_url(store_like)
+
+        if parsed.scheme == "memory":
+            # Create or get a ManagedMemoryStore
+            return ManagedMemoryStore(name=parsed.name, path=parsed.path, read_only=_read_only)
+        elif parsed.scheme == "file" or not parsed.scheme:
+            # Local filesystem path
+            return await make_store(Path(store_like), mode=mode, storage_options=storage_options)
+        else:
+            # Assume fsspec can handle it (s3://, gs://, http://, etc.)
             return FsspecStore.from_url(
                 store_like, storage_options=storage_options, read_only=_read_only
             )
-        else:
-            # Assume a filesystem path
-            return await make_store(Path(store_like), mode=mode, storage_options=storage_options)
 
     elif _has_fsspec and isinstance(store_like, FSMap):
         return FsspecStore.from_mapper(store_like, read_only=_read_only)
@@ -418,30 +415,25 @@ 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):
+        parsed = parse_store_url(store_like)
+        if parsed.scheme == "memory":
+            # Handle memory:// URLs specially - the path in the URL becomes part of the store
+            _read_only = mode == "r"
+            memory_store = ManagedMemoryStore(
+                name=parsed.name, path=parsed.path, read_only=_read_only
+            )
+            return await StorePath.open(memory_store, path=path_normalized, mode=mode)
+        else:
+            # Fall through to make_store for other URL types
+            store = await make_store(store_like, mode=mode, storage_options=storage_options)
+            return await StorePath.open(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)
 
 
-def _is_fsspec_uri(uri: str) -> bool:
-    """
-    Check if a URI looks like a non-local fsspec URI.
-
-    Examples
-    --------
-    ```python
-    from zarr.storage._common import _is_fsspec_uri
-    _is_fsspec_uri("s3://bucket")
-    # True
-    _is_fsspec_uri("my-directory")
-    # False
-    _is_fsspec_uri("local://my-directory")
-    # False
-    ```
-    """
-    return "://" in uri or ("::" in uri and "local://" not in uri)
-
-
 async def ensure_no_existing_node(
     store_path: StorePath,
     zarr_format: ZarrFormat,

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