fix zarr fsstore regression (#136)

* fix zarr fsstore regression

* fix pre-commit

* make mypy happy

Author: rabernat

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
 -   repo: https://github.com/pycqa/isort
-    rev: 5.6.4
+    rev: 5.12.0
     hooks:
       - id: isort
         args: ["--profile", "black", "--filter-files"]
@@ -9,12 +9,12 @@ repos:
     hooks:
     - id: black
       language_version: python3
--   repo: https://gitlab.com/pycqa/flake8
+-   repo: https://github.com/pycqa/flake8
     rev: 3.7.9
     hooks:
     - id: flake8
       language_version: python3
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.790
+    rev: v1.0.1
     hooks:
     -   id: mypy

rechunker/pipeline.py

@@ -1,6 +1,6 @@
 import itertools
 import math
-from typing import Any, Iterable, Iterator, Tuple, TypeVar
+from typing import Any, Iterable, Iterator, Tuple
 
 import dask
 import numpy as np
@@ -74,11 +74,8 @@ def specs_to_pipelines(specs: Iterable[CopySpec]) -> ParallelPipelines:
     return tuple((spec_to_pipeline(spec) for spec in specs))
 
 
-T = TypeVar("T")
-
-
 class CopySpecToPipelinesMixin(CopySpecExecutor):
-    def prepare_plan(self, specs: Iterable[CopySpec]) -> T:
+    def prepare_plan(self, specs: Iterable[CopySpec]):
         pipelines = specs_to_pipelines(specs)
         return self.pipelines_to_plan(pipelines)
 

tests/test_rechunk.py

@@ -9,6 +9,7 @@
 import numpy as np
 import pytest
 import zarr
+from zarr.storage import FSStore
 
 from rechunker import api
 
@@ -208,9 +209,9 @@ def test_rechunk_dataset(
     xarray = pytest.importorskip("xarray")
 
     if target_store.startswith("mapper"):
-        fsspec = pytest.importorskip("fsspec")
-        target_store = fsspec.get_mapper(str(tmp_path) + target_store)
-        temp_store = fsspec.get_mapper(str(tmp_path) + temp_store)
+        pytest.importorskip("fsspec")
+        target_store = FSStore(str(tmp_path) + target_store)
+        temp_store = FSStore(str(tmp_path) + temp_store)
     else:
         target_store = str(tmp_path / target_store)
         temp_store = str(tmp_path / temp_store)
@@ -445,16 +446,16 @@ def test_rechunk_dask_array(
 @pytest.mark.parametrize("temp_store", ["temp.zarr", "mapper.temp.zarr"])
 def test_rechunk_group(tmp_path, executor, source_store, target_store, temp_store):
     if source_store.startswith("mapper"):
-        fsspec = pytest.importorskip("fsspec")
-        store_source = fsspec.get_mapper(str(tmp_path) + source_store)
-        target_store = fsspec.get_mapper(str(tmp_path) + target_store)
-        temp_store = fsspec.get_mapper(str(tmp_path) + temp_store)
+        pytest.importorskip("fsspec")
+        store_source = FSStore(str(tmp_path) + source_store)
+        target_store = FSStore(str(tmp_path) + target_store)
+        temp_store = FSStore(str(tmp_path) + temp_store)
     else:
         store_source = str(tmp_path / source_store)
         target_store = str(tmp_path / target_store)
         temp_store = str(tmp_path / temp_store)
 
-    group = zarr.group(store_source)
+    group = zarr.group(store_source, overwrite=True)
     group.attrs["foo"] = "bar"
     # 800 byte chunks
     a = group.ones("a", shape=(5, 10, 20), chunks=(1, 10, 20), dtype="f4")