1515from abc import ABC , abstractmethod
1616from collections .abc import Sequence
1717from contextlib import nullcontext
18+ from tempfile import TemporaryDirectory
1819from typing import TYPE_CHECKING , Any
1920
2021import numpy as np
2122import torch
2223
23- from monai .utils import ensure_tuple_size , optional_import , require_pkg
24+ from monai .utils import ensure_tuple_size , get_package_version , optional_import , require_pkg , version_geq
2425
2526if TYPE_CHECKING :
2627 import zarr
@@ -233,7 +234,7 @@ def __init__(
233234 store : zarr .storage .Store | str = "merged.zarr" ,
234235 value_store : zarr .storage .Store | str | None = None ,
235236 count_store : zarr .storage .Store | str | None = None ,
236- compressor : str = "default" ,
237+ compressor : str | None = None ,
237238 value_compressor : str | None = None ,
238239 count_compressor : str | None = None ,
239240 chunks : Sequence [int ] | bool = True ,
@@ -246,8 +247,22 @@ def __init__(
246247 self .value_dtype = value_dtype
247248 self .count_dtype = count_dtype
248249 self .store = store
249- self .value_store = zarr .storage .TempStore () if value_store is None else value_store
250- self .count_store = zarr .storage .TempStore () if count_store is None else count_store
250+ self .tmpdir : TemporaryDirectory | None
251+ if version_geq (get_package_version ("zarr" ), "3.0.0" ):
252+ if value_store is None :
253+ self .tmpdir = TemporaryDirectory ()
254+ self .value_store = zarr .storage .LocalStore (self .tmpdir .name )
255+ else :
256+ self .value_store = value_store
257+ if count_store is None :
258+ self .tmpdir = TemporaryDirectory ()
259+ self .count_store = zarr .storage .LocalStore (self .tmpdir .name )
260+ else :
261+ self .count_store = count_store
262+ else :
263+ self .tmpdir = None
264+ self .value_store = zarr .storage .TempStore () if value_store is None else value_store
265+ self .count_store = zarr .storage .TempStore () if count_store is None else count_store
251266 self .chunks = chunks
252267 self .compressor = compressor
253268 self .value_compressor = value_compressor
0 commit comments