Merge pull request #3304 from tabedzki/main

Adding return type annotations

Author: alejoe91

src/spikeinterface/comparison/comparisontools.py

@@ -3,12 +3,10 @@
 """
 
 from __future__ import annotations
-
-
 import numpy as np
 
 
-def count_matching_events(times1, times2, delta=10):
+def count_matching_events(times1, times2: np.ndarray | list, delta: int = 10):
     """
     Counts matching events.
 
@@ -39,7 +37,7 @@ def count_matching_events(times1, times2, delta=10):
     return len(inds2) + 1
 
 
-def compute_agreement_score(num_matches, num1, num2):
+def compute_agreement_score(num_matches: int, num1: int, num2: int) -> float:
     """
     Computes agreement score.
 
@@ -232,8 +230,8 @@ def make_match_count_matrix(sorting1, sorting2, delta_frames, ensure_symmetry=Fa
         And the minimum of the two results is taken.
     Returns
     -------
-    matching_matrix : ndarray
-        A 2D numpy array of shape `(num_units_train1, num_units_train2)`. Each element `[i, j]` represents
+    matching_matrix : pd.DataFrame
+        A 2D pandas DataFrame of shape `(num_units_train1, num_units_train2)`. Each element `[i, j]` represents
         the count of matching spike pairs between unit `i` from `spike_frames_train1` and unit `j` from `spike_frames_train2`.
 
     Notes
@@ -348,7 +346,7 @@ def make_agreement_scores(sorting1, sorting2, delta_frames, ensure_symmetry=True
         And the minimum of the two results is taken.
     Returns
     -------
-    agreement_scores : array (float)
+    agreement_scores : pd.DataFrame
         The agreement score matrix.
     """
     import pandas as pd
@@ -408,9 +406,9 @@ def make_possible_match(agreement_scores, min_score):
 
     Returns
     -------
-    best_match_12 : pd.Series
+    best_match_12 : dict[NDArray]
 
-    best_match_21 : pd.Series
+    best_match_21 : dict[NDArray]
 
     """
     unit1_ids = np.array(agreement_scores.index)
@@ -433,7 +431,7 @@ def make_possible_match(agreement_scores, min_score):
     return possible_match_12, possible_match_21
 
 
-def make_best_match(agreement_scores, min_score):
+def make_best_match(agreement_scores, min_score) -> "tuple[pd.Series, pd.Series]":
     """
     Given an agreement matrix and a min_score threshold.
     return a dict a best match for each units independently of others.

src/spikeinterface/core/baseevent.py

@@ -137,7 +137,7 @@ class BaseEventSegment(BaseSegment):
     def __init__(self):
         BaseSegment.__init__(self)
 
-    def get_event_times(self, channel_id: int | str, start_time: float, end_time: float):
+    def get_event_times(self, channel_id: int | str, start_time: float, end_time: float) -> np.ndarray:
         """Returns event timestamps of a channel in seconds
         Parameters
         ----------

src/spikeinterface/core/baserecording.py

@@ -305,7 +305,7 @@ def get_traces(
         order: "C" | "F" | None = None,
         return_scaled: bool = False,
         cast_unsigned: bool = False,
-    ):
+    ) -> np.ndarray:
         """Returns traces from recording.
 
         Parameters

src/spikeinterface/core/core_tools.py

@@ -257,8 +257,7 @@ def set_value_in_extractor_dict(extractor_dict: dict, access_path: tuple, new_va
 
     Returns
     -------
-    dict
-        The modified dictionary
+    None
     """
 
     current = extractor_dict
@@ -439,7 +438,7 @@ def make_paths_relative(input_dict: dict, relative_folder: str | Path) -> dict:
     return output_dict
 
 
-def make_paths_absolute(input_dict, base_folder):
+def make_paths_absolute(input_dict, base_folder) -> dict:
     """
     Recursively transform a dict describing an BaseExtractor to make every path absolute given a base_folder.
 
@@ -632,7 +631,7 @@ def normal_pdf(x, mu: float = 0.0, sigma: float = 1.0):
     return 1 / (sigma * np.sqrt(2 * np.pi)) * np.exp(-((x - mu) ** 2) / (2 * sigma**2))
 
 
-def retrieve_importing_provenance(a_class):
+def retrieve_importing_provenance(a_class) -> dict:
     """
     Retrieve the import provenance of a class, including its import name (that consists of the class name and the module),
     the top-level module, and the module version.

src/spikeinterface/core/generate.py

@@ -33,7 +33,7 @@ def generate_recording(
     set_probe: bool | None = True,
     ndim: int | None = 2,
     seed: int | None = None,
-) -> BaseRecording:
+) -> NumpySorting:
     """
     Generate a lazy recording object.
     Useful for testing API and algos.

src/spikeinterface/core/old_api_utils.py

@@ -88,7 +88,7 @@ def get_unit_ids(self):
         """
         return list(self._unit_map.keys())
 
-    def get_unit_spike_train(self, unit_id, start_frame=None, end_frame=None):
+    def get_unit_spike_train(self, unit_id, start_frame=None, end_frame=None) -> np.ndarray:
         """This function extracts spike frames from the specified unit.
         It will return spike frames from within three ranges:
 
@@ -122,7 +122,7 @@ def get_unit_spike_train(self, unit_id, start_frame=None, end_frame=None):
             unit_id=self._unit_map[unit_id], segment_index=0, start_frame=start_frame, end_frame=end_frame
         )
 
-    def get_units_spike_train(self, unit_ids=None, start_frame=None, end_frame=None):
+    def get_units_spike_train(self, unit_ids=None, start_frame=None, end_frame=None) -> list[np.ndarray]:
         """This function extracts spike frames from the specified units.
 
         Parameters
@@ -137,7 +137,7 @@ def get_units_spike_train(self, unit_ids=None, start_frame=None, end_frame=None)
 
         Returns
         -------
-        spike_train: numpy.ndarray
+        spike_train: list[numpy.ndarray]
             An 2D array containing all the frames for each spike in the
             specified units given the range of start and end frames
         """
@@ -146,7 +146,7 @@ def get_units_spike_train(self, unit_ids=None, start_frame=None, end_frame=None)
         spike_trains = [self.get_unit_spike_train(uid, start_frame, end_frame) for uid in unit_ids]
         return spike_trains
 
-    def get_sampling_frequency(self):
+    def get_sampling_frequency(self) -> float:
         """
         It returns the sampling frequency.
 

src/spikeinterface/core/recording_tools.py

@@ -640,7 +640,7 @@ def get_noise_levels(
     method: Literal["mad", "std"] = "mad",
     force_recompute: bool = False,
     **random_chunk_kwargs,
-):
+) -> np.ndarray:
     """
     Estimate noise for each channel using MAD methods.
     You can use standard deviation with `method="std"`

src/spikeinterface/core/sortinganalyzer.py

@@ -44,7 +44,7 @@ def create_sorting_analyzer(
     return_scaled=True,
     overwrite=False,
     **sparsity_kwargs,
-):
+) -> "SortingAnalyzer":
     """
     Create a SortingAnalyzer by pairing a Sorting and the corresponding Recording.
 
@@ -145,7 +145,7 @@ def create_sorting_analyzer(
     return sorting_analyzer
 
 
-def load_sorting_analyzer(folder, load_extensions=True, format="auto", storage_options=None):
+def load_sorting_analyzer(folder, load_extensions=True, format="auto", storage_options=None) -> "SortingAnalyzer":
     """
     Load a SortingAnalyzer object from disk.
 
@@ -1123,7 +1123,7 @@ def get_num_units(self) -> int:
         return self.sorting.get_num_units()
 
     ## extensions zone
-    def compute(self, input, save=True, extension_params=None, verbose=False, **kwargs):
+    def compute(self, input, save=True, extension_params=None, verbose=False, **kwargs) -> "AnalyzerExtension | None":
         """
         Compute one extension or several extensiosn.
         Internally calls compute_one_extension() or compute_several_extensions() depending on the input type.
@@ -1190,7 +1190,7 @@ def compute(self, input, save=True, extension_params=None, verbose=False, **kwar
         else:
             raise ValueError("SortingAnalyzer.compute() need str, dict or list")
 
-    def compute_one_extension(self, extension_name, save=True, verbose=False, **kwargs):
+    def compute_one_extension(self, extension_name, save=True, verbose=False, **kwargs) -> "AnalyzerExtension":
         """
         Compute one extension.
 

src/spikeinterface/core/waveform_tools.py

@@ -791,7 +791,7 @@ def estimate_templates_with_accumulator(
     return_std: bool = False,
     verbose: bool = False,
     **job_kwargs,
-):
+) -> np.ndarray:
     """
     This is a fast implementation to compute template averages and standard deviations.
     This is useful to estimate sparsity without the need to allocate large waveform buffers.

src/spikeinterface/core/waveforms_extractor_backwards_compatibility.py

@@ -373,7 +373,7 @@ def load_waveforms(
     with_recording: bool = True,
     sorting: Optional[BaseSorting] = None,
     output="MockWaveformExtractor",
-):
+) -> MockWaveformExtractor | SortingAnalyzer:
     """
     This read an old WaveformsExtactor folder (folder or zarr) and convert it into a SortingAnalyzer or MockWaveformExtractor.