Implement ChunkableMixin

src/spikeinterface/core/__init__.py

@@ -95,7 +95,7 @@
     get_best_job_kwargs,
     ensure_n_jobs,
     ensure_chunk_size,
-    ChunkRecordingExecutor,
+    ChunkExecutor,
     split_job_kwargs,
     fix_job_kwargs,
 )

src/spikeinterface/core/base.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from pathlib import Path
 import shutil
 from typing import Any
@@ -87,6 +85,8 @@ def __init__(self, main_ids: Sequence) -> None:
                 self._main_ids.dtype.kind in "uiSU"
             ), f"Main IDs can only be integers (signed/unsigned) or strings, not {self._main_ids.dtype}"
 
+        self._segments: "list[BaseSegment]" = []
+
         # dict at object level
         self._annotations = {}
 
@@ -142,11 +142,18 @@ def name(self, value):
             # we remove the annotation if it exists
             _ = self._annotations.pop("name", None)
 
+    @property
+    def segments(self) -> "list[BaseSegment]":
+        return self._segments
+
+    def add_segment(self, segment: "BaseSegment") -> None:
+        self._segments.append(segment)
+        segment.set_parent_extractor(self)
+
     def get_num_segments(self) -> int:
-        # This is implemented in BaseRecording or BaseSorting
-        raise NotImplementedError
+        return len(self._segments)
 
-    def get_parent(self) -> BaseExtractor | None:
+    def get_parent(self) -> "BaseExtractor | None":
         """Returns parent object if it exists, otherwise None"""
         return getattr(self, "_parent", None)
 
@@ -381,7 +388,7 @@ def delete_property(self, key) -> None:
 
     def copy_metadata(
         self,
-        other: BaseExtractor,
+        other: "BaseExtractor",
         only_main: bool = False,
         ids: Iterable | slice | None = None,
         skip_properties: Iterable[str] | None = None,
@@ -570,7 +577,7 @@ def to_dict(
         return dump_dict
 
     @staticmethod
-    def from_dict(dictionary: dict, base_folder: Path | str | None = None) -> BaseExtractor:
+    def from_dict(dictionary: dict, base_folder: Path | str | None = None) -> "BaseExtractor":
         """
         Instantiate extractor from dictionary
 
@@ -624,7 +631,7 @@ def save_metadata_to_folder(self, folder_metadata):
             values = self.get_property(key)
             np.save(prop_folder / (key + ".npy"), values)
 
-    def clone(self) -> BaseExtractor:
+    def clone(self) -> "BaseExtractor":
         """
         Clones an existing extractor into a new instance.
         """
@@ -816,7 +823,7 @@ def dump_to_pickle(
         file_path.write_bytes(pickle.dumps(dump_dict))
 
     @staticmethod
-    def load(file_or_folder_path: str | Path, base_folder: Path | str | bool | None = None) -> BaseExtractor:
+    def load(file_or_folder_path: str | Path, base_folder: Path | str | bool | None = None) -> "BaseExtractor":
         """
         Load extractor from file path (.json or .pkl)
 
@@ -839,7 +846,7 @@ def __reduce__(self):
         return (instance_constructor, intialization_args)
 
     @staticmethod
-    def load_from_folder(folder) -> BaseExtractor:
+    def load_from_folder(folder) -> "BaseExtractor":
         return BaseExtractor.load(folder)
 
     def _save(self, folder, **save_kwargs):
@@ -855,7 +862,7 @@ def _extra_metadata_to_folder(self, folder):
         # This implemented in BaseRecording for probe
         pass
 
-    def save(self, **kwargs) -> BaseExtractor:
+    def save(self, **kwargs) -> "BaseExtractor":
         """
         Save a SpikeInterface object.
 
@@ -891,7 +898,7 @@ def save(self, **kwargs) -> BaseExtractor:
 
     save.__doc__ = save.__doc__.format(_shared_job_kwargs_doc)
 
-    def save_to_memory(self, sharedmem=True, **save_kwargs) -> BaseExtractor:
+    def save_to_memory(self, sharedmem=True, **save_kwargs) -> "BaseExtractor":
         save_kwargs.pop("format", None)
 
         cached = self._save(format="memory", sharedmem=sharedmem, **save_kwargs)
@@ -1092,7 +1099,7 @@ def save_to_zarr(
         return cached
 
 
-def _load_extractor_from_dict(dic) -> BaseExtractor:
+def _load_extractor_from_dict(dic) -> "BaseExtractor":
     """
     Convert a dictionary into an instance of BaseExtractor or its subclass.