Add support for in-memory datasets (e.g. useful for testing) (#1310)

* Add support for in-memory datasets (e.g. useful for testing)

* Better error handling on dataset creation

* format, write changelog

Author: Tobias314

webknossos/Changelog.md

@@ -15,6 +15,7 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 ### Breaking Changes
 
 ### Added
+- Add support for in-memory zarr3 datasets (using the Tensorstore in-memory KvStore driver). [#1310](https://github.com/scalableminds/webknossos-libs/pull/1310)
 
 ### Changed
 - Added url parsing in webknossos download CLI command, to get webknossos url without environment variable. [#1299](https://github.com/scalableminds/webknossos-libs/pull/1299)

webknossos/tests/dataset/test_remote_dataset.py

@@ -0,0 +1,25 @@
+import numpy as np
+
+import webknossos as wk
+from tests.utils import TestTemporaryDirectoryNonLocal
+
+
+def test_create_dataset_remote_storage() -> None:
+    """Test creating a dataset with remote storage."""
+    # Create a temporary directory for the dataset
+    # with tempfile.TemporaryDirectory() as temp_dir:
+    with TestTemporaryDirectoryNonLocal() as temp_dir:
+        dataset = wk.Dataset(temp_dir / "ds", voxel_size=(10, 10, 10), exist_ok=True)
+        layer = dataset.add_layer(
+            "color",
+            wk.COLOR_CATEGORY,
+            data_format="zarr3",
+            bounding_box=wk.BoundingBox((0, 0, 0), (16, 16, 16)),
+        )
+        mag1 = layer.add_mag(1)
+        mag1.write(np.ones((16, 16, 16), dtype="uint8"))
+        ds = wk.Dataset.open(temp_dir / "ds")
+        read_data = ds.get_layer("color").get_mag(1).read()
+        assert read_data.shape == (1, 16, 16, 16)
+        assert read_data.dtype == np.uint8
+        assert np.all(read_data == 1)

webknossos/tests/utils.py

@@ -0,0 +1,16 @@
+import uuid
+from collections.abc import Generator
+from contextlib import contextmanager
+
+from upath import UPath
+
+
+@contextmanager
+def TestTemporaryDirectoryNonLocal() -> Generator[UPath, None, None]:
+    """Gives a temporary directory as UPath which does not use the "local" protocol (local file system).
+    Useful for testing functionality that uses non-local UPaths.
+    Currently implemented to use an in-memory file system. (no persistence across lifetime of the process)."""
+    random_prefix = str(uuid.uuid4())
+    temp_dir = UPath(f"memory:///{random_prefix}")
+    temp_dir.mkdir(parents=True, exist_ok=True)
+    yield temp_dir

webknossos/webknossos/dataset/_array.py

@@ -22,6 +22,8 @@
 from ..utils import is_fs_path
 from .data_format import DataFormat
 
+TS_CONTEXT = tensorstore.Context()
+
 
 def _is_power_of_two(num: int) -> bool:
     return num & (num - 1) == 0
@@ -442,6 +444,14 @@ def _make_kvstore(path: Path) -> str | dict[str, str | list[str]]:
             else:
                 kvstore_spec["aws_credentials"] = {"type": "default"}
             return kvstore_spec
+        elif isinstance(path, UPath) and path.protocol == "memory":
+            # use memory driver (in-memory file systems), e.g. useful for testing
+            # attention: this is not a persistent storage and it does not support
+            # multiprocessing since memory is not shared between processes
+            return {
+                "driver": "memory",
+                "path": path.path,
+            }
         else:
             return {
                 "driver": "file",
@@ -470,6 +480,7 @@ def _open(cls, path: Path) -> Self:
                 },
                 open=True,
                 create=False,
+                context=TS_CONTEXT,
             ).result()  # check that everything exists
             return cls(path, _array)
         except Exception as exc:
@@ -532,7 +543,8 @@ def resize(self, new_bbox: NDBoundingBox) -> None:
                 {
                     "driver": str(self.data_format),
                     "kvstore": self._make_kvstore(self._path),
-                }
+                },
+                context=TS_CONTEXT,
             ).result()
             if array.domain != current_array.domain:
                 raise RuntimeError(
@@ -632,7 +644,8 @@ def _array(self) -> tensorstore.TensorStore:
                     {
                         "driver": str(self.data_format),
                         "kvstore": self._make_kvstore(self._path),
-                    }
+                    },
+                    context=TS_CONTEXT,
                 ).result()
             except Exception as e:
                 raise ArrayException(
@@ -780,7 +793,8 @@ def create(cls, path: Path, array_info: ArrayInfo) -> "Zarr3Array":
                     ],
                 },
                 "create": True,
-            }
+            },
+            context=TS_CONTEXT,
         ).result()
         return cls(path, _array)
 
@@ -856,7 +870,8 @@ def create(cls, path: Path, array_info: ArrayInfo) -> "Zarr2Array":
                     "dimension_separator": "/",
                 },
                 "create": True,
-            }
+            },
+            context=TS_CONTEXT,
         ).result()
         return cls(path, _array)
 

webknossos/webknossos/dataset/dataset.py

@@ -352,20 +352,16 @@ def __init__(
         else:
             assert not read_only
 
-            dataset_path_exists = False
-            dataset_path_is_empty = False
-            try:
-                dataset_path_is_empty = next(self.path.iterdir(), None) is None
-                dataset_path_exists = True
-            except NotADirectoryError:
-                dataset_path_exists = True
-            except FileNotFoundError:
-                dataset_path_exists = False
-
-            if dataset_path_exists and not dataset_path_is_empty:
-                raise RuntimeError(
-                    f"Creation of Dataset at {self.path} failed, because a file or folder already exists at this path."
-                )
+            if self.path.exists():
+                if self.path.is_dir():
+                    if next(self.path.iterdir(), None) is not None:
+                        raise RuntimeError(
+                            f"Creation of Dataset at {self.path} failed, because a non-empty folder already exists at this path."
+                        )
+                else:
+                    raise NotADirectoryError(
+                        f"Creation of Dataset at {self.path} failed, because the given path already exists but is not a directory."
+                    )
             # Create directories on disk and write datasource-properties.json
             try:
                 self.path.mkdir(parents=True, exist_ok=True)