Skip to content

Commit e210f92

Browse files
authored
fix(storage): pickle the zarr.storage module by name (#4557)
* fix(storage): pickle the zarr.storage module by name `zarr.storage` swaps its class for a `VerboseModule` subclass to warn when `default_compressor` is set. Serializers that dispatch on the exact type, such as cloudpickle (Dask's default), had no reduce for the subclass and raised `TypeError: cannot pickle 'VerboseModule' object`. Give `VerboseModule` a `__reduce__` that re-imports the module by name, so `pickle` and `cloudpickle` both round-trip it. The deprecation warning is unchanged; removing it is left for a minor release. Refs #4029. Assisted-by: ClaudeCode:claude-opus-5-5 * chore: name the changelog fragment after the pull request Assisted-by: ClaudeCode:claude-opus-5-5
1 parent 957a5f6 commit e210f92

3 files changed

Lines changed: 17 additions & 0 deletions

File tree

‎changes/4557.bugfix.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
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‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
1+
import importlib
12
import sys
23
import warnings
4+
from collections.abc import Callable
35
from types import ModuleType
46
from typing import Any
57

@@ -41,5 +43,10 @@ def __setattr__(self, attr: str, value: Any) -> None:
4143
else:
4244
super().__setattr__(attr, value)
4345

46+
def __reduce__(self) -> tuple[Callable[[str], ModuleType], tuple[str]]:
47+
# Serializers such as cloudpickle dispatch on the exact type, so they have no
48+
# reduce for this subclass. Pickle the module by name instead.
49+
return importlib.import_module, (self.__name__,)
50+
4451

4552
sys.modules[__name__].__class__ = VerboseModule

‎tests/test_store/test_core.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import pickle
12
import tempfile
23
from collections.abc import Awaitable, Callable, Generator
34
from pathlib import Path
@@ -434,3 +435,11 @@ def test_different_open_mode(tmp_path: Path) -> None:
434435
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().",
435436
):
436437
zarr.open_array(store=zip_store, path="a", zarr_format=2, mode="r")
438+
439+
440+
def test_storage_module_pickles_by_name() -> None:
441+
"""
442+
Pickling `zarr.storage` re-imports the module, so objects that refer to it can be
443+
serialized, for example by Dask. See issue #4029.
444+
"""
445+
assert pickle.loads(pickle.dumps(zarr.storage)) is zarr.storage

0 commit comments

Comments
 (0)