Merge pull request #3130 from alejoe91/fix-drift-hybrid-serialiibility

Fix serializability in `save` function and `InjectDriftingTemplatesRecording`

Author: samuelgarcia

src/spikeinterface/core/base.py

@@ -7,7 +7,6 @@
 import weakref
 import json
 import pickle
-import os
 import random
 import string
 from packaging.version import parse
@@ -951,13 +950,14 @@ def save_to_folder(
         folder.mkdir(parents=True, exist_ok=False)
 
         # dump provenance
-        provenance_file = folder / f"provenance.json"
         if self.check_serializability("json"):
+            provenance_file = folder / f"provenance.json"
+            self.dump(provenance_file)
+        elif self.check_serializability("pickle"):
+            provenance_file = folder / f"provenance.pkl"
             self.dump(provenance_file)
         else:
-            provenance_file.write_text(
-                json.dumps({"warning": "the provenace is not json serializable!!!"}), encoding="utf8"
-            )
+            warnings.warn("The extractor is not serializable to file. The provenance will not be saved.")
 
         self.save_metadata_to_folder(folder)
 
@@ -1024,7 +1024,6 @@ def save_to_zarr(
         cached: ZarrExtractor
             Saved copy of the extractor.
         """
-        import zarr
         from .zarrextractors import read_zarr
 
         save_kwargs.pop("format", None)

src/spikeinterface/core/generate.py

@@ -1803,6 +1803,8 @@ def __init__(
             )
             self.add_recording_segment(recording_segment)
 
+        # to discuss: maybe we could set json serializability to False always
+        # because templates could be large!
         if not sorting.check_serializability("json"):
             self._serializability["json"] = False
         if parent_recording is not None:

src/spikeinterface/generation/drift_tools.py

@@ -458,6 +458,9 @@ def __init__(
 
         self.set_probe(drifting_templates.probe, in_place=True)
 
+        # templates are too large, we don't serialize them to JSON
+        self._serializability["json"] = False
+
         self._kwargs = {
             "sorting": sorting,
             "drifting_templates": drifting_templates,