Repository: SpikeInterface/spikeinterface

Issue:
[feat-core]: refactor API for saving to zarr
Instead of `zarr_path`, use `folder` + `format="zarr"`

Discussion:

Patch:
diff --git a/spikeinterface/core/base.py b/spikeinterface/core/base.py
index 53d0b926a..82081e240 100644
--- a/spikeinterface/core/base.py
+++ b/spikeinterface/core/base.py
@@ -1,6 +1,6 @@
 from pathlib import Path
 import importlib
-from copy import deepcopy
+import warnings
 import weakref
 import json
 import pickle
@@ -8,6 +8,7 @@ import os
 import random
 import string
 from packaging.version import parse
+from copy import deepcopy
 
 import numpy as np
 
@@ -709,8 +710,8 @@ class BaseExtractor:
 
         return cached
 
-    def save_to_zarr(self, name=None, zarr_path=None, storage_options=None, 
-                     channel_chunk_size=None, verbose=True, **save_kwargs):
+    def save_to_zarr(self, name=None, folder=None, storage_options=None, 
+                     channel_chunk_size=None, verbose=True, zarr_path=None, **save_kwargs):
         """
         Save extractor to zarr.
 
@@ -724,20 +725,33 @@ class BaseExtractor:
         name: str or None
             Name of the subfolder in get_global_tmp_folder()
             If 'name' is given, 'folder' must be None.
-        zarr_path: str, Path, or None
-            Name of the zarr folder (.zarr).
+        folder: str, Path, or None
+            The folder used to save the zarr output. If the folder does not have a '.zarr' suffix, 
+            it will be automatically appended.
         storage_options: dict or None
             Storage options for zarr `store`. E.g., if "s3://" or "gcs://" they can provide authentication methods, etc.
             For cloud storage locations, this should not be None (in case of default values, use an empty dict)
         channel_chunk_size: int or None
             Channels per chunk. Default None (chunking in time only)
+        verbose: bool
+            If True (default), the output is verbose.
+        zarr_path: str, Path, or None
+            (Deprecated) Name of the zarr folder (.zarr). 
         
         Returns
         -------
-        cached: saved copy of the extractor.
+        cached: ZarrRecordingExtractor
+            Saved copy of the extractor.
         """
         import zarr
-        if zarr_path is None:
+
+        if zarr_path is not None:
+            warnings.warn("The 'zarr_path' argument is deprecated. "
+                          "Use 'folder' instead",
+                          DeprecationWarning, stacklevel=2)
+            folder = zarr_path
+
+        if folder is None:
             cache_folder = get_global_tmp_folder()
             if name is None:
                 name = ''.join(random.choices(
@@ -752,12 +766,13 @@ class BaseExtractor:
                         print(f'Use zarr_path={zarr_path}')
         else:
             if storage_options is None:
-                if isinstance(zarr_path, str):
-                    zarr_path_init = zarr_path
-                    zarr_path = Path(zarr_path)
-                else:
-                    zarr_path_init = str(zarr_path)
+                folder = Path(folder)
+                if folder.suffix != ".zarr":
+                    folder = folder.parent / f"{folder.stem}.zarr"
+                zarr_path = folder
+                zarr_path_init = str(zarr_path)
             else:
+                zarr_path = folder
                 zarr_path_init = zarr_path
 
         if isinstance(zarr_path, Path):


Test patch:
diff --git a/spikeinterface/core/tests/test_baserecording.py b/spikeinterface/core/tests/test_baserecording.py
index 57f72b8b3..dea84c950 100644
--- a/spikeinterface/core/tests/test_baserecording.py
+++ b/spikeinterface/core/tests/test_baserecording.py
@@ -255,18 +255,18 @@ def test_BaseRecording():
 
     # Test save to zarr
     compressor = get_default_zarr_compressor()
-    rec_zarr = rec2.save(format="zarr", zarr_path=cache_folder / "recording.zarr",
+    rec_zarr = rec2.save(format="zarr", folder=cache_folder / "recording",
                          compressor=compressor)
     check_recordings_equal(rec2, rec_zarr, return_scaled=False)
 
-    rec_zarr2 = rec2.save(format="zarr", zarr_path=cache_folder / "recording_channel_chunk.zarr",
+    rec_zarr2 = rec2.save(format="zarr", folder=cache_folder / "recording_channel_chunk",
                           compressor=compressor, channel_chunk_size=2)
     check_recordings_equal(rec2, rec_zarr2, return_scaled=False)
 
     # test cast unsigned
-    rec_u = rec_uint16.save(format="zarr", zarr_path=cache_folder / "rec_u.zarr")
+    rec_u = rec_uint16.save(format="zarr", folder=cache_folder / "rec_u")
     rec_u.get_dtype() == 'uint16'
-    rec_i = rec_uint16.save(format="zarr", zarr_path=cache_folder / "rec_i.zarr", dtype="int16")
+    rec_i = rec_uint16.save(format="zarr", folder=cache_folder / "rec_i", dtype="int16")
     rec_i.get_dtype() == 'int16'
     assert np.allclose(rec_u.get_traces(cast_unsigned=False).astype("float") - (2**15),
                        rec_i.get_traces().astype("float"))