fix ty checks

Author: will

src/funlib/persistence/arrays/datasets.py

@@ -1,6 +1,6 @@
 import logging
 from collections.abc import Sequence
-from typing import Any, Optional, Union
+from typing import Any, Literal, Optional, Union
 
 import numpy as np
 import zarr
@@ -14,7 +14,7 @@
 
 def open_ds(
     store,
-    mode: str = "r",
+    mode: Literal["r", "r+", "a", "w", "w-"] = "r",
     metadata_format: Optional[MetaDataFormat] = None,
     offset: Optional[Sequence[int]] = None,
     voxel_size: Optional[Sequence[int]] = None,
@@ -131,7 +131,7 @@ def prepare_ds(
     types: Optional[Sequence[str]] = None,
     chunk_shape: Optional[Sequence[int]] = None,
     dtype: DTypeLike = np.float32,
-    mode: str = "a",
+    mode: Literal["r", "r+", "a", "w", "w-"] = "a",
     custom_metadata: dict[str, Any] | None = None,
     **kwargs,
 ) -> Array:

src/funlib/persistence/arrays/metadata.py

@@ -1,5 +1,5 @@
 import warnings
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from pathlib import Path
 from typing import Any, Optional
 
@@ -268,7 +268,7 @@ class Config:
         extra = "forbid"
 
     def fetch(
-        self, data: dict[str | int, Any], key: str
+        self, data: Mapping[str, Any], key: str
     ) -> Sequence[str | int | None] | None:
         """
         Given a dictionary of attributes from e.g. zarr.open(...).attrs, fetch the value
@@ -292,7 +292,7 @@ def fetch(
         keys = key.split("/")
 
         def recurse(
-            data: dict[str | int, Any] | list[Any], keys: list[str]
+            data: Any, keys: list[str]
         ) -> Sequence[str | int | None] | str | int | None:
             current_key: str | int
             current_key, *keys = keys
@@ -304,7 +304,7 @@ def recurse(
             # base case
             if len(keys) == 0:
                 # this key returns the data we want
-                if isinstance(data, (dict, zarr.core.attributes.Attributes)):
+                if isinstance(data, Mapping):
                     return data.get(str(current_key), None)
                 elif isinstance(data, list):
                     assert isinstance(current_key, int), current_key
@@ -336,7 +336,7 @@ def recurse(
     def parse(
         self,
         shape: Sequence[int],
-        data: dict[str | int, Any],
+        data: Mapping[str, Any],
         offset: Optional[Sequence[int]] = None,
         voxel_size: Optional[Sequence[int]] = None,
         axis_names: Optional[Sequence[str]] = None,

src/funlib/persistence/arrays/ome_datasets.py

@@ -63,8 +63,8 @@ def open_ome_ds(
 
     metadata = MetaData(
         shape=dataset.shape,
-        offset=offset,  # type: ignore[arg-type]
-        voxel_size=scale,  # type: ignore[arg-type]
+        offset=offset,
+        voxel_size=scale,
         axis_names=axis_names,
         units=units,
         types=types,