diff --git a/changes/4557.bugfix.md b/changes/4557.bugfix.md new file mode 100644 index 0000000000..b2c76744ef --- /dev/null +++ b/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`. diff --git a/src/zarr/storage/__init__.py b/src/zarr/storage/__init__.py index f1bd1724af..9a2580c066 100644 --- a/src/zarr/storage/__init__.py +++ b/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 diff --git a/tests/test_store/test_core.py b/tests/test_store/test_core.py index daf62e3843..0653b77a57 100644 --- a/tests/test_store/test_core.py +++ b/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