Merge branch 'main' into implement_canonical_structs

Author: d-v-b

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

@@ -23,8 +23,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 _join_paths, normalize_path, parse_store_url
 
 _has_fsspec = importlib.util.find_spec("fsspec")
 if _has_fsspec:
@@ -36,18 +36,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.
@@ -267,10 +255,10 @@ def delete_sync(self) -> None:
 
     def __truediv__(self, other: str) -> StorePath:
         """Combine this store path with another path"""
-        return self.__class__(self.store, _dereference_path(self.path, other))
+        return self.__class__(self.store, _join_paths([self.path, other]))
 
     def __str__(self) -> str:
-        return _dereference_path(str(self.store), self.path)
+        return _join_paths([str(self.store), self.path])
 
     def __repr__(self) -> str:
         return f"StorePath({self.store.__class__.__name__}, '{self}')"
@@ -342,14 +330,17 @@ 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.",
-        )
+    # Parse URL early so we can reuse the result for both validation and routing
+    parsed = parse_store_url(store_like) if isinstance(store_like, str) else None
+
+    # Check if storage_options is valid for this store_like
+    if storage_options is not None:
+        is_fsspec_uri = parsed is not None and 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"
@@ -377,15 +368,18 @@ async def make_store(
         # Create a new LocalStore
         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):
+    elif isinstance(store_like, str) and parsed is not None:
+        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 — use parsed.path to strip the file:// scheme
+            return await make_store(Path(parsed.path), 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)
@@ -460,25 +454,6 @@ async def make_store_path(
         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._utils import _join_paths
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterable
@@ -282,7 +282,7 @@ async def get(
         # docstring inherited
         if not self._is_open:
             await self._open()
-        path = _dereference_path(self.path, key)
+        path = _join_paths([self.path, key])
 
         try:
             if byte_range is None:
@@ -329,7 +329,7 @@ async def set(
             raise TypeError(
                 f"FsspecStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
             )
-        path = _dereference_path(self.path, key)
+        path = _join_paths([self.path, key])
         # write data
         if byte_range:
             raise NotImplementedError
@@ -338,7 +338,7 @@ async def set(
     async def delete(self, key: str) -> None:
         # docstring inherited
         self._check_writable()
-        path = _dereference_path(self.path, key)
+        path = _join_paths([self.path, key])
         try:
             await self.fs._rm(path)
         except FileNotFoundError:
@@ -354,14 +354,14 @@ async def delete_dir(self, prefix: str) -> None:
             )
         self._check_writable()
 
-        path_to_delete = _dereference_path(self.path, prefix)
+        path_to_delete = _join_paths([self.path, prefix])
 
         with suppress(*self.allowed_exceptions):
             await self.fs._rm(path_to_delete, recursive=True)
 
     async def exists(self, key: str) -> bool:
         # docstring inherited
-        path = _dereference_path(self.path, key)
+        path = _join_paths([self.path, key])
         exists: bool = await self.fs._exists(path)
         return exists
 
@@ -378,7 +378,7 @@ async def get_partial_values(
             starts: list[int | None] = []
             stops: list[int | None] = []
             for key, byte_range in key_ranges:
-                paths.append(_dereference_path(self.path, key))
+                paths.append(_join_paths([self.path, key]))
                 if byte_range is None:
                     starts.append(None)
                     stops.append(None)
@@ -429,7 +429,7 @@ async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
             yield onefile.removeprefix(f"{self.path}/")
 
     async def getsize(self, key: str) -> int:
-        path = _dereference_path(self.path, key)
+        path = _join_paths([self.path, key])
         info = await self.fs._info(path)
 
         size = info.get("size")