fix(storage): pickle the zarr.storage module by name

changes/4557.bugfix.md

@@ -0,0 +1 @@
+The `zarr.storage` module can now be serialized with `pickle` and `cloudpickle`. Previously, code that referenced the module (for example a Dask task that used `zarr.storage.ObjectStore`) failed with `TypeError: cannot pickle 'VerboseModule' object`.

src/zarr/storage/__init__.py

@@ -1,5 +1,7 @@
+import importlib
 import sys
 import warnings
+from collections.abc import Callable
 from types import ModuleType
 from typing import Any
 
@@ -41,5 +43,10 @@ def __setattr__(self, attr: str, value: Any) -> None:
         else:
             super().__setattr__(attr, value)
 
+    def __reduce__(self) -> tuple[Callable[[str], ModuleType], tuple[str]]:
+        # Serializers such as cloudpickle dispatch on the exact type, so they have no
+        # reduce for this subclass. Pickle the module by name instead.
+        return importlib.import_module, (self.__name__,)
+
 
 sys.modules[__name__].__class__ = VerboseModule

tests/test_store/test_core.py

@@ -1,3 +1,4 @@
+import pickle
 import tempfile
 from collections.abc import Awaitable, Callable, Generator
 from pathlib import Path
@@ -434,3 +435,11 @@ def test_different_open_mode(tmp_path: Path) -> None:
         match="Store is not read-only but mode is 'r'. Unable to create a read-only copy of the store. Please use a read-only store or a storage class that implements .with_read_only().",
     ):
         zarr.open_array(store=zip_store, path="a", zarr_format=2, mode="r")
+
+
+def test_storage_module_pickles_by_name() -> None:
+    """
+    Pickling `zarr.storage` re-imports the module, so objects that refer to it can be
+    serialized, for example by Dask. See issue #4029.
+    """
+    assert pickle.loads(pickle.dumps(zarr.storage)) is zarr.storage