ngio 0.5.0__py3-none-any.whl → 0.5.0a1__py3-none-any.whl

ngio/utils/_zarr_utils.py

@@ -1,67 +1,58 @@
 """Common utilities for working with Zarr groups in consistent ways."""
 
-import json
-import warnings
 from pathlib import Path
-from typing import Literal, TypeAlias
+from typing import Literal
 
-import dask.array as da
 import fsspec
 import zarr
 from filelock import BaseFileLock, FileLock
-from pydantic_zarr.v2 import ArraySpec as AnyArraySpecV2
-from pydantic_zarr.v3 import ArraySpec as AnyArraySpecV3
 from zarr.abc.store import Store
+from zarr.core.array import CompressorLike
 from zarr.errors import ContainsGroupError
-from zarr.storage import FsspecStore, LocalStore, MemoryStore, ZipStore
+from zarr.storage import FsspecStore, LocalStore, MemoryStore
 
-from ngio.utils._cache import NgioCache
-from ngio.utils._errors import (
-    NgioFileExistsError,
-    NgioFileNotFoundError,
-    NgioValueError,
-)
+from ngio.utils import NgioFileExistsError, NgioFileNotFoundError, NgioValueError
+from ngio.utils._errors import NgioError
 
 AccessModeLiteral = Literal["r", "r+", "w", "w-", "a"]
 # StoreLike is more restrictive than it could be
 # but to make sure we can handle the store correctly
 # we need to be more restrictive
-NgioSupportedStore: TypeAlias = (
-    str | Path | fsspec.mapping.FSMap | FsspecStore | MemoryStore | dict | LocalStore
+NgioSupportedStore = (
+    str | Path | fsspec.mapping.FSMap | FsspecStore | MemoryStore | LocalStore
 )
-GenericStore: TypeAlias = NgioSupportedStore | Store
-StoreOrGroup: TypeAlias = NgioSupportedStore | zarr.Group
+GenericStore = Store | NgioSupportedStore
+StoreOrGroup = GenericStore | zarr.Group
 
 
 def _check_store(store) -> NgioSupportedStore:
     """Check the store and return a valid store."""
-    if not isinstance(store, NgioSupportedStore):
-        warnings.warn(
-            f"Store type {type(store)} is not explicitly supported. "
-            f"Supported types are: {NgioSupportedStore}. "
-            "Proceeding, but this may lead to unexpected behavior.",
-            UserWarning,
-            stacklevel=2,
-        )
-    return store
+    if isinstance(store, NgioSupportedStore):
+        return store
 
+    raise NotImplementedError(
+        f"Store type {type(store)} is not supported. "
+        f"Supported types are: {NgioSupportedStore}"
+    )
 
-def _check_group(
-    group: zarr.Group, mode: AccessModeLiteral | None = None
-) -> zarr.Group:
+
+def _check_group(group: zarr.Group, mode: AccessModeLiteral) -> zarr.Group:
     """Check the group and return a valid group."""
-    if group.read_only and mode not in [None, "r"]:
-        raise NgioValueError(f"The group is read only. Cannot open in mode {mode}.")
+    if group.read_only and mode in ["w", "w-"]:
+        raise NgioValueError(
+            "The group is read only. Cannot open in write mode ['w', 'w-']"
+        )
 
     if mode == "r" and not group.read_only:
         # let's make sure we don't accidentally write to the group
         group = zarr.open_group(store=group.store, path=group.path, mode="r")
+
     return group
 
 
 def open_group_wrapper(
     store: StoreOrGroup,
-    mode: AccessModeLiteral | None = None,
+    mode: AccessModeLiteral,
     zarr_format: Literal[2, 3] | None = None,
 ) -> zarr.Group:
     """Wrapper around zarr.open_group with some additional checks.
@@ -81,7 +72,6 @@ def open_group_wrapper(
 
     try:
         _check_store(store)
-        mode = mode if mode is not None else "a"
         group = zarr.open_group(store=store, mode=mode, zarr_format=zarr_format)
 
     except FileExistsError as e:
@@ -108,32 +98,68 @@ class ZarrGroupHandler:
         store: StoreOrGroup,
         zarr_format: Literal[2, 3] | None = None,
         cache: bool = False,
-        mode: AccessModeLiteral | None = None,
+        mode: AccessModeLiteral = "a",
+        parallel_safe: bool = False,
+        parent: "ZarrGroupHandler | None" = None,
     ):
         """Initialize the handler.
 
         Args:
             store (StoreOrGroup): The Zarr store or group containing the image data.
-            zarr_format (int | None): The Zarr format version to use.
+            meta_mode (str): The mode of the metadata handler.
+            zarr_format (int): The Zarr format version to use.
             cache (bool): Whether to cache the metadata.
-            mode (str | None): The mode of the store.
+            mode (str): The mode of the store.
+            parallel_safe (bool): If True, the handler will create a lock file to make
+                that can be used to make the handler parallel safe.
+                Be aware that the lock needs to be used manually.
+            parent (ZarrGroupHandler | None): The parent handler.
         """
-        if mode not in ["r", "r+", "w", "w-", "a", None]:
+        if mode not in ["r", "r+", "w", "w-", "a"]:
             raise NgioValueError(f"Mode {mode} is not supported.")
 
+        if parallel_safe and cache:
+            raise NgioValueError(
+                "The cache and parallel_safe options are mutually exclusive."
+                "If you want to use the lock mechanism, you should not use the cache."
+            )
+
         group = open_group_wrapper(store=store, mode=mode, zarr_format=zarr_format)
+        _store = group.store
+
+        # Make sure the cache is set in the attrs
+        # in the same way as the cache in the handler
+
+        ## TODO
+        # Figure out how to handle the cache in the new zarr version
+        # group.attrs.cache = cache
+
+        if parallel_safe:
+            if not isinstance(_store, LocalStore):
+                raise NgioValueError(
+                    "The store needs to be a LocalStore to use the lock mechanism. "
+                    f"Instead, got {_store.__class__.__name__}."
+                )
+
+            store_path = _store.root / group.path
+            self._lock_path = store_path.with_suffix(".lock")
+            self._lock = FileLock(self._lock_path, timeout=10)
+
+        else:
+            self._lock_path = None
+            self._lock = None
+
         self._group = group
+        self._mode = mode
         self.use_cache = cache
-
-        self._group_cache: NgioCache[zarr.Group] = NgioCache(use_cache=cache)
-        self._array_cache: NgioCache[zarr.Array] = NgioCache(use_cache=cache)
-        self._handlers_cache: NgioCache[ZarrGroupHandler] = NgioCache(use_cache=cache)
-        self._lock: tuple[Path, BaseFileLock] | None = None
+        self._parallel_safe = parallel_safe
+        self._cache = {}
+        self._parent = parent
 
     def __repr__(self) -> str:
         """Return a string representation of the handler."""
         return (
-            f"ZarrGroupHandler(full_url={self.full_url}, read_only={self.read_only}, "
+            f"ZarrGroupHandler(full_url={self.full_url}, mode={self.mode}, "
             f"cache={self.use_cache}"
         )
 
@@ -147,17 +173,8 @@ class ZarrGroupHandler:
         """Return the store path."""
         if isinstance(self.store, LocalStore):
             return (self.store.root / self.group.path).as_posix()
-        elif isinstance(self.store, FsspecStore):
-            return f"{self.store.path}/{self.group.path}"
-        elif isinstance(self.store, ZipStore):
-            return (self.store.path / self.group.path).as_posix()
-        elif isinstance(self.store, MemoryStore):
-            return None
-        warnings.warn(
-            f"Cannot determine full URL for store type {type(self.store)}. ",
-            UserWarning,
-            stacklevel=2,
-        )
+        if isinstance(self.store, FsspecStore):
+            return self.store.fs.map.root_path
         return None
 
     @property
@@ -166,55 +183,35 @@ class ZarrGroupHandler:
         return self._group.metadata.zarr_format
 
     @property
-    def read_only(self) -> bool:
-        """Return whether the group is read only."""
-        return self._group.read_only
-
-    def _create_lock(self) -> tuple[Path, BaseFileLock]:
-        """Create the lock."""
-        if self._lock is not None:
-            return self._lock
-
-        if self.use_cache is True:
-            raise NgioValueError(
-                "Lock mechanism is not compatible with caching. "
-                "Please set cache=False to use the lock mechanism."
-            )
-
-        if not isinstance(self.store, LocalStore):
-            raise NgioValueError(
-                "The store needs to be a LocalStore to use the lock mechanism. "
-                f"Instead, got {self.store.__class__.__name__}."
-            )
-
-        store_path = Path(self.store.root) / self.group.path
-        _lock_path = store_path.with_suffix(".lock")
-        _lock = FileLock(_lock_path, timeout=10)
-        return _lock_path, _lock
+    def mode(self) -> AccessModeLiteral:
+        """Return the mode of the group."""
+        return self._mode  # type: ignore
 
     @property
     def lock(self) -> BaseFileLock:
         """Return the lock."""
         if self._lock is None:
-            self._lock = self._create_lock()
-        return self._lock[1]
+            raise NgioValueError(
+                "The handler is not parallel safe. "
+                "Reopen the handler with parallel_safe=True."
+            )
+        return self._lock
 
     @property
-    def lock_path(self) -> Path:
-        """Return the lock path."""
-        if self._lock is None:
-            self._lock = self._create_lock()
-        return self._lock[0]
+    def parent(self) -> "ZarrGroupHandler | None":
+        """Return the parent handler."""
+        return self._parent
 
     def remove_lock(self) -> None:
         """Return the lock."""
-        if self._lock is None:
+        if self._lock is None or self._lock_path is None:
             return None
 
-        lock_path, lock = self._lock
-        if lock_path.exists() and lock.lock_counter == 0:
+        lock_path = Path(self._lock_path)
+        if lock_path.exists() and self._lock.lock_counter == 0:
             lock_path.unlink()
             self._lock = None
+            self._lock_path = None
             return None
 
         raise NgioValueError("The lock is still in use. Cannot remove it.")
@@ -225,7 +222,10 @@ class ZarrGroupHandler:
         This is useful when the group has been modified
         outside of the handler.
         """
-        mode = "r" if self.read_only else "r+"
+        if self.mode == "r":
+            mode = "r"
+        else:
+            mode = "r+"
         return zarr.open_group(
             store=self._group.store,
             path=self._group.path,
@@ -233,57 +233,72 @@ class ZarrGroupHandler:
             zarr_format=self._group.metadata.zarr_format,
         )
 
-    def reopen_handler(self) -> "ZarrGroupHandler":
-        """Reopen the handler.
-
-        This is useful when the group has been modified
-        outside of the handler.
-        """
-        mode = "r" if self.read_only else "r+"
-        group = self.reopen_group()
-        return ZarrGroupHandler(
-            store=group,
-            zarr_format=group.metadata.zarr_format,
-            cache=self.use_cache,
-            mode=mode,
-        )
-
-    def clean_cache(self) -> None:
-        """Clear the cached metadata."""
-        group = self.reopen_group()
-        self.__init__(
-            store=group,
-            zarr_format=group.metadata.zarr_format,
-            cache=self.use_cache,
-            mode="r" if self.read_only else "r+",
-        )
-
     @property
     def group(self) -> zarr.Group:
         """Return the group."""
-        if self.use_cache is False:
-            # If we are not using cache, we need to reopen the group
+        if self._parallel_safe:
+            # If we are parallel safe, we need to reopen the group
             # to make sure that the attributes are up to date
             return self.reopen_group()
         return self._group
 
+    def add_to_cache(self, key: str, value: object) -> None:
+        """Add an object to the cache."""
+        if not self.use_cache:
+            return None
+        self._cache[key] = value
+
+    def get_from_cache(self, key: str) -> object | None:
+        """Get an object from the cache."""
+        if not self.use_cache:
+            return None
+        return self._cache.get(key, None)
+
+    def clean_cache(self) -> None:
+        """Clear the cached metadata."""
+        self._cache = {}
+
     def load_attrs(self) -> dict:
         """Load the attributes of the group."""
-        return self.reopen_group().attrs.asdict()
+        attrs = self.get_from_cache("attrs")
+        if attrs is not None and isinstance(attrs, dict):
+            return attrs
 
-    def write_attrs(self, attrs: dict, overwrite: bool = False) -> None:
+        attrs = dict(self.group.attrs)
+
+        self.add_to_cache("attrs", attrs)
+        return attrs
+
+    def _write_attrs(self, attrs: dict, overwrite: bool = False) -> None:
         """Write the metadata to the store."""
-        # Maybe we should use the lock here
-        if self.read_only:
+        if self.group.read_only:
             raise NgioValueError("The group is read only. Cannot write metadata.")
-        group = self.reopen_group()
+
+        # we need to invalidate the current attrs cache
+        self.add_to_cache("attrs", None)
         if overwrite:
-            group.attrs.clear()
-        group.attrs.update(attrs)
+            self.group.attrs.clear()
+
+        self.group.attrs.update(attrs)
+
+    def write_attrs(self, attrs: dict, overwrite: bool = False) -> None:
+        """Write the metadata to the store."""
+        # Maybe we should use the lock here
+        self._write_attrs(attrs, overwrite)
+
+    def _obj_get(self, path: str):
+        """Get a group from the group."""
+        group_or_array = self.get_from_cache(path)
+        if group_or_array is not None:
+            return group_or_array
+
+        group_or_array = self.group.get(path, None)
+        self.add_to_cache(path, group_or_array)
+        return group_or_array
 
     def create_group(self, path: str, overwrite: bool = False) -> zarr.Group:
         """Create a group in the group."""
-        if self.group.read_only:
+        if self.mode == "r":
             raise NgioValueError("Cannot create a group in read only mode.")
 
         try:
@@ -293,7 +308,7 @@ class ZarrGroupHandler:
                 f"A Zarr group already exists at {path}, "
                 "consider setting overwrite=True."
             ) from e
-        self._group_cache.set(path, group, overwrite=overwrite)
+        self.add_to_cache(path, group)
         return group
 
     def get_group(
@@ -319,87 +334,138 @@ class ZarrGroupHandler:
         if overwrite:
             return self.create_group(path, overwrite=overwrite)
 
-        group = self._group_cache.get(path)
+        group = self._obj_get(path)
         if isinstance(group, zarr.Group):
             return group
 
-        group = self.group.get(path, default=None)
-        if isinstance(group, zarr.Group):
-            self._group_cache.set(path, group, overwrite=overwrite)
-            return group
-
-        if isinstance(group, zarr.Array):
-            raise NgioValueError(f"The object at {path} is not a group, but an array.")
+        if group is not None:
+            raise NgioValueError(
+                f"The object at {path} is not a group, but a {type(group)}"
+            )
 
         if not create_mode:
             raise NgioFileNotFoundError(f"No group found at {path}")
         group = self.create_group(path)
-        self._group_cache.set(path, group, overwrite=overwrite)
         return group
 
+    def safe_get_group(
+        self, path: str, create_mode: bool = False
+    ) -> tuple[bool, zarr.Group | NgioError]:
+        """Get a group from the group.
+
+        Args:
+            path (str): The path to the group.
+            create_mode (bool): If True, create the group if it does not exist.
+
+        Returns:
+            zarr.Group | None: The Zarr group or None if it does not exist
+                or an error occurs.
+
+        """
+        try:
+            return True, self.get_group(path, create_mode)
+        except NgioError as e:
+            return False, e
+
     def get_array(self, path: str) -> zarr.Array:
         """Get an array from the group."""
-        array = self._array_cache.get(path)
-        if isinstance(array, zarr.Array):
-            return array
-        array = self.group.get(path, default=None)
-        if isinstance(array, zarr.Array):
-            self._array_cache.set(path, array)
-            return array
-
-        if isinstance(array, zarr.Group):
-            raise NgioValueError(f"The object at {path} is not an array, but a group.")
-        raise NgioFileNotFoundError(f"No array found at {path}")
-
-    def get_handler(
+        array = self._obj_get(path)
+        if array is None:
+            raise NgioFileNotFoundError(f"No array found at {path}")
+        if not isinstance(array, zarr.Array):
+            raise NgioValueError(
+                f"The object at {path} is not an array, but a {type(array)}"
+            )
+        return array
+
+    def create_array(
+        self,
+        path: str,
+        shape: tuple[int, ...],
+        dtype: str,
+        chunks: tuple[int, ...] | Literal["auto"] = "auto",
+        compressors: CompressorLike = "auto",
+        separator: Literal[".", "/"] = "/",
+        overwrite: bool = False,
+    ) -> zarr.Array:
+        if self.mode == "r":
+            raise NgioValueError("Cannot create an array in read only mode.")
+
+        if self.zarr_format == 2:
+            chunks_encoding = {
+                "name": "v2",
+                "separator": separator,
+            }
+        else:
+            chunks_encoding = {
+                "name": "default",
+                "separator": separator,
+            }
+
+        try:
+            return self.group.create_array(
+                name=path,
+                shape=shape,
+                dtype=dtype,
+                chunks=chunks,
+                chunk_key_encoding=chunks_encoding,
+                overwrite=overwrite,
+                compressors=compressors,
+            )
+        except ContainsGroupError as e:
+            raise NgioFileExistsError(
+                f"A Zarr array already exists at {path}, "
+                "consider setting overwrite=True."
+            ) from e
+        except Exception as e:
+            raise NgioValueError(f"Error creating array at {path}") from e
+
+    def derive_handler(
         self,
         path: str,
-        create_mode: bool = True,
         overwrite: bool = False,
     ) -> "ZarrGroupHandler":
-        """Get a new handler for a group in the current handler group.
+        """Derive a new handler from the current handler.
 
         Args:
             path (str): The path to the group.
-            create_mode (bool): If True, create the group if it does not exist.
             overwrite (bool): If True, overwrite the group if it exists.
         """
-        handler = self._handlers_cache.get(path)
-        if handler is not None:
-            return handler
-        group = self.get_group(path, create_mode=create_mode, overwrite=overwrite)
-        mode = "r" if group.read_only else "r+"
-        handler = ZarrGroupHandler(
-            store=group, zarr_format=self.zarr_format, cache=self.use_cache, mode=mode
+        group = self.get_group(path, create_mode=True, overwrite=overwrite)
+        return ZarrGroupHandler(
+            store=group,
+            zarr_format=self.zarr_format,
+            cache=self.use_cache,
+            mode=self.mode,
+            parallel_safe=self._parallel_safe,
+            parent=self,
         )
-        self._handlers_cache.set(path, handler)
-        return handler
-
-    @property
-    def is_listable(self) -> bool:
-        return is_group_listable(self.group)
 
-    def delete_group(self, path: str) -> None:
-        """Delete a group from the current group.
-
-        Args:
-            path (str): The path to the group to delete.
-        """
-        if self.group.read_only:
-            raise NgioValueError("Cannot delete a group in read only mode.")
-        self.group.__delitem__(path)
-        self._group_cache._cache.pop(path, None)
-        self._handlers_cache._cache.pop(path, None)
-
-    def delete_self(self) -> None:
-        """Delete the current group."""
-        if self.group.read_only:
-            raise NgioValueError("Cannot delete a group in read only mode.")
-        self.group.__delitem__("/")
+    def safe_derive_handler(
+        self,
+        path: str,
+        overwrite: bool = False,
+    ) -> tuple[bool, "ZarrGroupHandler | NgioError"]:
+        """Derive a new handler from the current handler."""
+        try:
+            return True, self.derive_handler(path, overwrite=overwrite)
+        except NgioError as e:
+            return False, e
 
-    def copy_group(self, dest_group: zarr.Group):
+    def copy_handler(self, handler: "ZarrGroupHandler") -> None:
         """Copy the group to a new store."""
-        copy_group(self.group, dest_group)
+        _, n_skipped, _ = zarr.copy_store(
+            source=self.group.store,
+            dest=handler.group.store,
+            source_path=self.group.path,
+            dest_path=handler.group.path,
+            if_exists="replace",
+        )
+        if n_skipped > 0:
+            raise NgioValueError(
+                f"Error copying group to {handler.full_url}, "
+                f"#{n_skipped} files where skipped."
+            )
 
 
 def find_dimension_separator(array: zarr.Array) -> Literal[".", "/"]:
@@ -418,116 +484,8 @@ def find_dimension_separator(array: zarr.Array) -> Literal[".", "/"]:
     else:
         separator = array.metadata.chunk_key_encoding
         if not isinstance(separator, DefaultChunkKeyEncoding):
-            raise NgioValueError(
+            raise ValueError(
                 "Only DefaultChunkKeyEncoding is supported in this example."
             )
         separator = separator.separator
     return separator
-
-
-def is_group_listable(group: zarr.Group) -> bool:
-    """Check if a Zarr group is listable.
-
-    A group is considered listable if it contains at least one array or subgroup.
-
-    Args:
-        group (zarr.Group): The Zarr group to check.
-
-    Returns:
-        bool: True if the group is listable, False otherwise.
-    """
-    if not group.store.supports_listing:
-        # If the store does not support listing
-        # then for sure it is not listable
-        return False
-    try:
-        next(group.keys())
-        return True
-    except StopIteration:
-        # Group is listable but empty
-        return True
-    except Exception as _:
-        # Some stores may raise errors when listing
-        # consider those not listable
-        return False
-
-
-def _make_sync_fs(fs: fsspec.AbstractFileSystem) -> fsspec.AbstractFileSystem:
-    fs_dict = json.loads(fs.to_json())
-    fs_dict["asynchronous"] = False
-    return fsspec.AbstractFileSystem.from_json(json.dumps(fs_dict))
-
-
-def _get_mapper(store: LocalStore | FsspecStore, path: str):
-    if isinstance(store, LocalStore):
-        fs = fsspec.filesystem("file")
-        full_path = (store.root / path).as_posix()
-    else:
-        fs = _make_sync_fs(store.fs)
-        full_path = f"{store.path}/{path}"
-    return fs.get_mapper(full_path)
-
-
-def _fsspec_copy(
-    src_fs: LocalStore | FsspecStore,
-    src_path: str,
-    dest_fs: LocalStore | FsspecStore,
-    dest_path: str,
-):
-    src_mapper = _get_mapper(src_fs, src_path)
-    dest_mapper = _get_mapper(dest_fs, dest_path)
-    for key in src_mapper.keys():
-        dest_mapper[key] = src_mapper[key]
-
-
-def _zarr_python_copy(src_group: zarr.Group, dest_group: zarr.Group):
-    # Copy attributes
-    dest_group.attrs.put(src_group.attrs.asdict())
-    # Copy arrays
-    for name, array in src_group.arrays():
-        if array.metadata.zarr_format == 2:
-            spec = AnyArraySpecV2.from_zarr(array)
-        else:
-            spec = AnyArraySpecV3.from_zarr(array)
-        dst = spec.to_zarr(
-            store=dest_group.store,
-            path=f"{dest_group.path}/{name}",
-            overwrite=True,
-        )
-        if array.ndim > 0:
-            dask_array = da.from_zarr(array)
-            da.to_zarr(dask_array, dst, overwrite=False)
-    # Copy subgroups
-    for name, subgroup in src_group.groups():
-        dest_subgroup = dest_group.create_group(name, overwrite=True)
-        _zarr_python_copy(subgroup, dest_subgroup)
-
-
-def copy_group(
-    src_group: zarr.Group, dest_group: zarr.Group, suppress_warnings: bool = False
-):
-    if src_group.metadata.zarr_format != dest_group.metadata.zarr_format:
-        raise NgioValueError(
-            "Different Zarr format versions between source and destination, "
-            "cannot copy."
-        )
-
-    if not is_group_listable(src_group):
-        raise NgioValueError("Source group is not listable, cannot copy.")
-
-    if dest_group.read_only:
-        raise NgioValueError("Destination group is read only, cannot copy.")
-    if isinstance(src_group.store, LocalStore | FsspecStore) and isinstance(
-        dest_group.store, LocalStore | FsspecStore
-    ):
-        _fsspec_copy(src_group.store, src_group.path, dest_group.store, dest_group.path)
-        return
-    if not suppress_warnings:
-        warnings.warn(
-            "Fsspec copy not possible, falling back to Zarr Python API for the copy. "
-            "This will preserve some tabular data non-zarr native (parquet, and csv), "
-            "and it will be slower for large datasets.",
-            UserWarning,
-            stacklevel=2,
-        )
-    _zarr_python_copy(src_group, dest_group)