Make obstore the default for cloud stores, if it is installed (#914)

Add test for use_obstore

Author: tomwhite

cubed/storage/stores/zarr_python_v3.py

@@ -4,6 +4,12 @@
 import zarr
 
 from cubed.types import T_DType, T_RegularChunks, T_Shape, T_Store
+from cubed.utils import is_cloud_storage_path
+
+try:
+    import obstore
+except ImportError:
+    obstore = None  # type: ignore
 
 zarr.config.set(
     {
@@ -46,21 +52,35 @@ def open_zarr_v3_array(
     path: Optional[str] = None,
     **kwargs,
 ):
-    # use obstore if requested
+    # use obstore if explicitly requested, or if library is installed and store is a cloud store
     storage_options = kwargs.pop("storage_options", None)
-    if storage_options is not None and storage_options.get("use_obstore", False):
-        import obstore as obs
+    obstore_requested = storage_options is not None and storage_options.get(
+        "use_obstore", False
+    )
+    obstore_installed = obstore is not None
+    if obstore_requested and not obstore_installed:
+        raise RuntimeError(
+            "obstore was requested with 'use_obstore=True' but it is not installed"
+        )
+    use_obstore = obstore_requested or (
+        obstore_installed
+        and isinstance(store, (str, Path))
+        and is_cloud_storage_path(store)
+    )
+    if use_obstore:
         from zarr.storage import ObjectStore
 
         if isinstance(store, str):
             if "://" not in store:
                 p = Path(store)
-                store = ObjectStore(obs.store.from_url(p.as_uri(), mkdir=True))
+                store = ObjectStore(obstore.store.from_url(p.as_uri(), mkdir=True))
             else:
-                store = ObjectStore(obs.store.from_url(store))
+                store = ObjectStore(obstore.store.from_url(store))
         elif isinstance(store, Path):
             p = store
-            store = ObjectStore(obs.store.from_url(p.as_uri(), mkdir=True))
+            store = ObjectStore(obstore.store.from_url(p.as_uri(), mkdir=True))
+        else:
+            raise RuntimeError("Store must be a string or `Path` object for obstore")
 
     if isinstance(chunks, int):
         chunks = (chunks,)

cubed/tests/storage/test_zarr.py

@@ -1,20 +1,36 @@
 import pytest
 from zarr.codecs import BloscCodec, BytesCodec, ZstdCodec
+from zarr.storage import LocalStore, ObjectStore
 
 from cubed import config
 from cubed.storage.store import open_storage_array
 from cubed.storage.zarr import lazy_zarr_array
 
 
-def test_lazy_zarr_array(tmp_path):
+@pytest.mark.parametrize("use_obstore", [None, False, True])
+def test_lazy_zarr_array(tmp_path, use_obstore):
     zarr_path = tmp_path / "lazy.zarr"
-    arr = lazy_zarr_array(zarr_path, shape=(3, 3), dtype=int, chunks=(2, 2))
+    if use_obstore is None:
+        kwargs = {}
+    else:
+        kwargs = dict(storage_options=dict(use_obstore=use_obstore))
+    arr = lazy_zarr_array(
+        zarr_path,
+        shape=(3, 3),
+        dtype=int,
+        chunks=(2, 2),
+        **kwargs,
+    )
 
     assert not zarr_path.exists()
     with pytest.raises((FileNotFoundError, TypeError, ValueError)):
         arr.open()
 
-    arr.create()
+    za = arr.create()
+    if use_obstore:
+        assert isinstance(za.store, ObjectStore)
+    else:
+        assert isinstance(za.store, LocalStore)
     assert zarr_path.exists()
     arr.open()
 

docs/user-guide/storage.md

@@ -38,7 +38,7 @@ To use the [`zarrs-python`](https://github.com/zarrs/zarrs-python) Rust implemen
 export CUBED_STORAGE_NAME=zarrs-python
 ```
 
-For `zarr-python` v3 only, you can use Zarr's [`ObjectStore`](https://zarr.readthedocs.io/en/main/api/zarr/storage/index.html#zarr.storage.ObjectStore) which uses the Rust-based [`obstore`](https://developmentseed.org/obstore/latest/), by setting the `storage_options.use_obstore` [configuration](../configuration.md) option to `True`, as illustrated in this YAML file:
+For `zarr-python`, if [`obstore`](https://developmentseed.org/obstore/latest/) is installed it will be used by default for cloud object stores. To enable it for other stores (e.g. the local filesystem), set the `storage_options.use_obstore` [configuration](../configuration.md) option to `True`, as illustrated in this YAML file:
 
 ```yaml
 spec: