Chec to_dict/load and improve warning msg

SpikeInterface · Jan 31, 2024 · 8ef06a9 · 8ef06a9
1 parent d7dedb5
commit 8ef06a9
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/src/spikeinterface/extractors/nwbextractors.py b/src/spikeinterface/extractors/nwbextractors.py
@@ -608,8 +608,8 @@ def __init__(
 
         if storage_options is not None and stream_mode == "zarr":
             warnings.warn(
-                "The `storage_options` parameter will not be propagated to the `kwargs` for security reasons. "
-                "The extractor will not be serializable."
+                "The `storage_options` parameter will not be propagated to JSON or pickle files for security reasons, "
+                "so the extractor will not be JSON/pickle serializable. Only in-memory mode will be available."
             )
             # not serializable if storage_options is provided
             self._serializability["json"] = False
@@ -622,6 +622,7 @@ def __init__(
             "samples_for_rate_estimation": samples_for_rate_estimation,
             "stream_mode": stream_mode,
             "load_channel_properties": load_channel_properties,
+            "storage_options": storage_options,
             "cache": cache,
             "stream_cache_path": stream_cache_path,
             "file": file,
@@ -1039,8 +1040,8 @@ def __init__(
 
         if storage_options is not None and stream_mode == "zarr":
             warnings.warn(
-                "The `storage_options` parameter will not be propagated to the `kwargs` for security reasons. "
-                "The extractor will not be serializable."
+                "The `storage_options` parameter will not be propagated to JSON or pickle files for security reasons, "
+                "so the extractor will not be JSON/pickle serializable. Only in-memory mode will be available."
             )
             # not serializable if storage_options is provided
             self._serializability["json"] = False
@@ -1054,6 +1055,7 @@ def __init__(
             "cache": cache,
             "stream_mode": stream_mode,
             "stream_cache_path": stream_cache_path,
+            "storage_options": storage_options,
             "load_unit_properties": load_unit_properties,
             "t_start": self.t_start,
         }

diff --git a/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py b/src/spikeinterface/extractors/tests/test_nwbextractors_streaming.py
@@ -1,11 +1,13 @@
 from pathlib import Path
 import pickle
+from tabnanny import check
 
 import pytest
 import numpy as np
 import h5py
-from spikeinterface.core.testing import check_recordings_equal
 
+from spikeinterface import load_extractor
+from spikeinterface.core.testing import check_recordings_equal
 from spikeinterface.core.testing import check_recordings_equal, check_sortings_equal
 from spikeinterface.extractors import NwbRecordingExtractor, NwbSortingExtractor
 
@@ -288,6 +290,10 @@ def test_sorting_s3_nwb_zarr(tmp_path):
     assert not sorting.check_serializability("json")
     assert not sorting.check_serializability("pickle")
 
+    # test to/from dict
+    sorting_loaded = load_extractor(sorting.to_dict())
+    check_sortings_equal(sorting, sorting_loaded)
+
 
 if __name__ == "__main__":
     tmp_path = Path("tmp")