Docstring to sorting analyzer functions (#991)

Co-authored-by: Cody Baker <[email protected]>
catalystneuro · Aug 8, 2024 · ef656a0 · ef656a0
1 parent fc6ae26
commit ef656a0
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 10 deletions.
diff --git a/src/neuroconv/tools/spikeinterface/spikeinterface.py b/src/neuroconv/tools/spikeinterface/spikeinterface.py
@@ -1271,7 +1271,7 @@ def add_units_table(
 def add_sorting(
     sorting: BaseSorting,
     nwbfile: Optional[pynwb.NWBFile] = None,
-    unit_ids: Optional[List[Union[str, int]]] = None,
+    unit_ids: Optional[Union[List[str], List[int]]] = None,
     property_descriptions: Optional[dict] = None,
     skip_properties: Optional[List[str]] = None,
     skip_features: Optional[List[str]] = None,
@@ -1439,15 +1439,21 @@ def add_sorting_analyzer(
     nwbfile: Optional[pynwb.NWBFile] = None,
     metadata: Optional[dict] = None,
     recording: Optional[BaseRecording] = None,
-    unit_ids: Optional[List[Union[str, int]]] = None,
+    unit_ids: Optional[Union[List[str], List[int]]] = None,
     skip_properties: Optional[List[str]] = None,
     property_descriptions: Optional[dict] = None,
     write_as: Literal["units", "processing"] = "units",
     units_name: str = "units",
     units_description: str = "Autogenerated by neuroconv.",
 ):
     """
-    Primary method for writing object to an NWBFile.
+    Convenience function to write directly a sorting analyzer object to an nwbfile.
+
+    The function adds the data of the recording and the sorting plus the following information from the sorting analyzer:
+    - quality metrics
+    - template mean and std
+    - template metrics
+
 
     Parameters
     ----------
@@ -1546,7 +1552,7 @@ def write_sorting_analyzer(
     overwrite: bool = False,
     recording: Optional[BaseRecording] = None,
     verbose: bool = True,
-    unit_ids: Optional[List[Union[str, int]]] = None,
+    unit_ids: Optional[Union[List[str], List[int]]] = None,
     write_electrical_series: bool = False,
     add_electrical_series_kwargs: Optional[dict] = None,
     skip_properties: Optional[List[str]] = None,
@@ -1556,7 +1562,12 @@ def write_sorting_analyzer(
     units_description: str = "Autogenerated by neuroconv.",
 ):
     """
-    Primary method for writing object to an NWBFile.
+    Convenience function to write directly a sorting analyzer object to an nwbfile.
+
+    The function adds the data of the recording and the sorting plus the following information from the sorting analyzer:
+    - quality metrics
+    - template mean and std
+    - template metrics
 
     Parameters
     ----------

diff --git a/tests/test_ecephys/test_tools_spikeinterface.py b/tests/test_ecephys/test_tools_spikeinterface.py
@@ -1276,16 +1276,11 @@ def test_adding_doubled_ragged_arrays(self):
             np.testing.assert_array_equal(value, expected_values[i])
 
 
-from packaging.version import Version
-
 from neuroconv.tools import get_package_version
 
 spike_interface_version = get_package_version("spikeinterface")
 
 
-@unittest.skipIf(
-    spike_interface_version >= Version("0.100"), reason="WaveformExtractor not available in spikeinterface"
-)
 class TestWriteSortingAnalyzer(TestCase):
     @classmethod
     def setUpClass(cls):