Remove Zarr pin (#667)

* Remove Zarr pin

* Update type hints for Zarr Python v3

* Escape pip package spec

Author: tomwhite

.github/workflows/zarr-v3-tests.yml renamed to .github/workflows/zarr-v2-tests.yml

@@ -1,4 +1,4 @@
-name: Zarr v3 Tests
+name: Zarr v2 Tests
 
 on:
   push:
@@ -44,7 +44,7 @@ jobs:
       - name: Install
         run: |
           python -m pip install -e .[test,lithops]
-          python -m pip install -U git+https://github.com/zarr-developers/zarr-python.git
+          python -m pip install -U 'zarr<3'
 
       - name: Run tests
         env:

cubed/primitive/blockwise.py

@@ -16,7 +16,7 @@
     numpy_array_to_backend_array,
 )
 from cubed.runtime.types import CubedPipeline
-from cubed.storage.zarr import T_ZarrArray, lazy_zarr_array
+from cubed.storage.zarr import LazyZarrArray, T_ZarrArray, lazy_zarr_array
 from cubed.types import T_Chunks, T_DType, T_Shape, T_Store
 from cubed.utils import array_memory, chunk_memory, get_item, map_nested
 from cubed.utils import numblocks as compute_numblocks
@@ -321,7 +321,7 @@ def general_blockwise(
 
     write_proxies = []
     output_chunk_memory = 0
-    target_array = []
+    target_arrays = []
 
     numblocks0 = None
     for i, target_store in enumerate(target_stores):
@@ -335,6 +335,7 @@ def general_blockwise(
                 raise ValueError(
                     f"All outputs must have matching number of blocks in each dimension. Chunks specified: {chunkss}"
                 )
+        ta: Union[zarr.Array, LazyZarrArray]
         if isinstance(target_store, zarr.Array):
             ta = target_store
         else:
@@ -347,7 +348,7 @@ def general_blockwise(
                 storage_options=storage_options,
                 compressor=compressor,
             )
-        target_array.append(ta)
+        target_arrays.append(ta)
 
         write_proxies.append(CubedArrayProxy(ta, chunksize))
 
@@ -356,9 +357,6 @@ def general_blockwise(
             output_chunk_memory, array_memory(dtypes[i], chunksize) * 2
         )
 
-    if len(target_array) == 1:
-        target_array = target_array[0]
-
     spec = BlockwiseSpec(
         key_function,
         func_with_kwargs,
@@ -404,7 +402,7 @@ def general_blockwise(
     return PrimitiveOperation(
         pipeline=pipeline,
         source_array_names=array_names,
-        target_array=target_array,
+        target_array=target_arrays[0] if len(target_arrays) == 1 else target_arrays,
         projected_mem=projected_mem,
         allowed_mem=allowed_mem,
         reserved_mem=reserved_mem,

cubed/primitive/rechunk.py

@@ -1,7 +1,7 @@
 import itertools
 import math
 from math import ceil, prod
-from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Tuple
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple
 
 import numpy as np
 
@@ -68,10 +68,9 @@ def rechunk(
         storage_options=storage_options,
     )
 
-    intermediate = int_proxy.array
     target = write_proxy.array
 
-    if intermediate is None:
+    if int_proxy is None:
         copy_spec = CubedCopySpec(read_proxy, write_proxy)
         num_tasks = total_chunks(write_proxy.array.shape, write_proxy.chunks)
         return [
@@ -88,6 +87,7 @@ def rechunk(
 
     else:
         # break spec into two if there's an intermediate
+        intermediate = int_proxy.array
         copy_spec1 = CubedCopySpec(read_proxy, int_proxy)
         num_tasks = total_chunks(copy_spec1.write.array.shape, copy_spec1.write.chunks)
         op1 = spec_to_primitive_op(
@@ -123,7 +123,7 @@ def _setup_array_rechunk(
     target_store: T_Store,
     temp_store: Optional[T_Store] = None,
     storage_options: Optional[Dict[str, Any]] = None,
-) -> Tuple[CubedArrayProxy, CubedArrayProxy, CubedArrayProxy]:
+) -> Tuple[CubedArrayProxy, Optional[CubedArrayProxy], CubedArrayProxy]:
     shape = source_array.shape
     source_chunks = source_array.chunks
     dtype = source_array.dtype
@@ -162,7 +162,9 @@ def _setup_array_rechunk(
         )
 
     read_proxy = CubedArrayProxy(source_array, read_chunks)
-    int_proxy = CubedArrayProxy(int_array, int_chunks)
+    int_proxy = (
+        CubedArrayProxy(int_array, int_chunks) if int_array is not None else None
+    )
     write_proxy = CubedArrayProxy(target_array, write_chunks)
     return read_proxy, int_proxy, write_proxy
 
@@ -200,7 +202,7 @@ def __iter__(self):
         return chunk_keys(self.shape, self.chunks)
 
 
-def copy_read_to_write(chunk_key: Sequence[slice], *, config: CubedCopySpec) -> None:
+def copy_read_to_write(chunk_key: Tuple[slice], *, config: CubedCopySpec) -> None:
     # workaround limitation of lithops.utils.verify_args
     if isinstance(chunk_key, list):
         chunk_key = tuple(chunk_key)

cubed/storage/backend.py

@@ -33,7 +33,7 @@ def open_backend_array(
     storage_name = backend_storage_name()
 
     if storage_name == "zarr-python":
-        from cubed.storage.backends.zarr_python import open_zarr_array
+        from cubed.storage.backends.zarr_python import open_zarr_array  # type: ignore
 
         open_func = open_zarr_array
 

cubed/storage/backends/zarr_python.py

@@ -1,3 +1,4 @@
+# type: ignore  # Zarr Python v2 types are different to v3
 from typing import Optional, Union
 
 import zarr

pyproject.toml

@@ -36,7 +36,7 @@ dependencies = [
     "psutil",
     "tenacity",
     "toolz",
-    "zarr < 3",
+    "zarr",
 ]
 
 [project.optional-dependencies]