diff --git a/src/spikeinterface/extractors/nwbextractors.py b/src/spikeinterface/extractors/nwbextractors.py index 0d1946d39a..5c6f2ea6ce 100644 --- a/src/spikeinterface/extractors/nwbextractors.py +++ b/src/spikeinterface/extractors/nwbextractors.py @@ -137,6 +137,25 @@ def read_nwbfile( if file_path is None and file is None: raise ValueError("Provide either file_path or file") + if ( + backend == "zarr" + and file is None + and _get_zarr_format(file_path, stream_mode, storage_options) == 2 + and _has_zarr_python_3() + ): + from hdmf_zarr import NWBZarrV2IO + + # hdmf-zarr>=0.14 reads zarr v2 files with a separate read-only IO. Files written by hdmf-zarr<0.14 + # pickle their object references (e.g. the "group" column of the electrodes table), so pickle is allowed + io = NWBZarrV2IO( + path=str(file_path), + mode="r", + load_namespaces=True, + allow_pickle=True, + storage_options=storage_options, + ) + return io.read() + open_file = read_file_from_backend( file_path=file_path, file=file, @@ -303,19 +322,45 @@ def _get_backend_from_local_file(file_path: str | Path) -> str: else: raise RuntimeError(f"{file_path} is not a valid HDF5 file!") elif file_path.is_dir(): - try: - import zarr - - # `mode` is keyword-only in zarr v3, and its groups are not context managers - zarr.open(file_path, mode="r") - backend = "zarr" - except: + # The root metadata file is checked instead of opening the folder, because zarr>=3 cannot open + # the consolidated metadata of the zarr v2 files written by hdmf-zarr<0.14 + if _get_zarr_format(file_path) is None: raise RuntimeError(f"{file_path} is not a valid Zarr folder!") + backend = "zarr" else: raise RuntimeError(f"File {file_path} is not an existing file or folder!") return backend +def _get_zarr_format(file_path: str | Path, stream_mode=None, storage_options: dict | None = None) -> int | None: + """ + Return the zarr format (2 or 3) of a zarr folder, local or streamed with stream_mode="zarr", from its + root metadata file, or None if the folder is not a zarr group. + """ + if stream_mode == "zarr": + import fsspec + + file_system, root = fsspec.core.url_to_fs(str(file_path), **(storage_options or {})) + root = root.rstrip("/") + exists = lambda name: file_system.exists(f"{root}/{name}") + elif stream_mode is None: + folder_path = Path(file_path) + exists = lambda name: (folder_path / name).is_file() + else: + return None + if exists("zarr.json"): + return 3 + if exists(".zgroup"): + return 2 + return None + + +def _has_zarr_python_3() -> bool: + import zarr + + return int(zarr.__version__.split(".")[0]) >= 3 + + def _find_neurodata_type_from_backend(group, path="", result=None, neurodata_type="ElectricalSeries", backend="hdf5"): """ Recursively searches for groups with the specified neurodata_type hdf5 or zarr object, @@ -402,6 +447,30 @@ def _storage_backend(file_path=None, file=None, stream_mode=None): return _get_backend_from_local_file(file_path) return "zarr" if stream_mode == "zarr" else "hdf5" + @staticmethod + def _reading_method(use_pynwb, file_path=None, file=None, stream_mode=None, storage_options=None): + # With zarr-python 3, the in-house zarr reader only reads zarr v3 files. Zarr v2 files (written by + # hdmf-zarr<0.14) are read with pynwb, where hdmf-zarr keeps the backward compatibility code for them. + if use_pynwb: + return "use_pynwb" + backend = _NWBReader._storage_backend(file_path, file, stream_mode) + if ( + backend == "zarr" + and _has_zarr_python_3() + and _get_zarr_format(file_path, stream_mode, storage_options) == 2 + ): + if not HAVE_PYNWB or importlib.util.find_spec("hdmf_zarr") is None: + raise ImportError( + "To read zarr v2 NWB files with zarr-python 3, install pynwb and hdmf-zarr>=0.14, which requires " + 'Python 3.12 or newer: \n\n pip install pynwb "hdmf-zarr>=0.14"\n\n' + ) + warnings.warn( + f"{file_path} is a zarr v2 NWB file, which is read with pynwb. This can be slower than " + "the default reader, which only reads zarr v3 files." + ) + return "use_pynwb" + return f"use_{backend}" + def _open_handle(self, *, file_path, file, stream_mode, cache, stream_cache_path, storage_options=None): # Open the file handle: a pynwb NWBFile, or a raw h5py.File / zarr group. if self.reading_method == "use_pynwb": @@ -1002,10 +1071,7 @@ def __init__( if use_pynwb and not HAVE_PYNWB: raise ImportError(self.installation_mesg) - if use_pynwb: - reading_method = "use_pynwb" # the reader detects hdf5 vs zarr itself for pynwb - else: - reading_method = f"use_{_NWBReader._storage_backend(file_path, file, self.stream_mode)}" + reading_method = _NWBReader._reading_method(use_pynwb, file_path, file, self.stream_mode, storage_options) self._reader = _NWBReader( reading_method=reading_method, @@ -1317,10 +1383,9 @@ def __init__( if use_pynwb and not HAVE_PYNWB: raise ImportError(self.installation_mesg) - if use_pynwb: - reading_method = "use_pynwb" - else: - reading_method = f"use_{_NWBReader._storage_backend(file_path, stream_mode=stream_mode)}" + reading_method = _NWBReader._reading_method( + use_pynwb, file_path, stream_mode=stream_mode, storage_options=storage_options + ) self._reader = _NWBReader( reading_method=reading_method, @@ -1725,10 +1790,7 @@ def __init__( if use_pynwb and not HAVE_PYNWB: raise ImportError(self.installation_mesg) - if use_pynwb: - reading_method = "use_pynwb" - else: - reading_method = f"use_{_NWBReader._storage_backend(file_path, file, stream_mode)}" + reading_method = _NWBReader._reading_method(use_pynwb, file_path, file, stream_mode, storage_options) self._reader = _NWBReader(reading_method=reading_method, timeseries_path=timeseries_path) self._reader.load_timeseries( diff --git a/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py b/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py index 9a7b2c0c67..161040cbaf 100644 --- a/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py +++ b/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py @@ -3,18 +3,12 @@ import pytest import numpy as np -import zarr from spikeinterface import load from spikeinterface.core.testing import check_recordings_equal from spikeinterface.core.testing import check_recordings_equal, check_sortings_equal from spikeinterface.extractors.extractor_classes import NwbRecordingExtractor, NwbSortingExtractor -# TODO: remove once streaming zarr is supported with zarr>=3 -skip_zarr_3 = pytest.mark.skipif( - int(zarr.__version__.split(".")[0]) >= 3, reason="Streaming zarr is not supported yet with zarr>=3" -) - @pytest.mark.streaming_extractors @pytest.mark.parametrize("cache", [True, False]) # Test with and without cache @@ -191,7 +185,6 @@ def test_sorting_s3_nwb_remfile(tmp_path): @pytest.mark.streaming_extractors -@skip_zarr_3 def test_sorting_s3_nwb_zarr(tmp_path): file_path = ( "s3://aind-open-data/ecephys_625749_2022-08-03_15-15-06_nwb_2023-05-16_16-34-55/"