Merge pull request #3287 from h-mayorquin/typing_to_set_property_and_annotations

Refactor `set_property` in base

Author: samuelgarcia

src/spikeinterface/core/base.py

@@ -164,7 +164,7 @@ def id_to_index(self, id) -> int:
     def annotate(self, **new_annotations) -> None:
         self._annotations.update(new_annotations)
 
-    def set_annotation(self, annotation_key, value: Any, overwrite=False) -> None:
+    def set_annotation(self, annotation_key: str, value: Any, overwrite=False) -> None:
         """This function adds an entry to the annotations dictionary.
 
         Parameters
@@ -192,7 +192,7 @@ def get_preferred_mp_context(self):
         """
         return self._preferred_mp_context
 
-    def get_annotation(self, key, copy: bool = True) -> Any:
+    def get_annotation(self, key: str, copy: bool = True) -> Any:
         """
         Get a annotation.
         Return a copy by default
@@ -205,7 +205,13 @@ def get_annotation(self, key, copy: bool = True) -> Any:
     def get_annotation_keys(self) -> List:
         return list(self._annotations.keys())
 
-    def set_property(self, key, values: Sequence, ids: Optional[Sequence] = None, missing_value: Any = None) -> None:
+    def set_property(
+        self,
+        key,
+        values: list | np.ndarray | tuple,
+        ids: list | np.ndarray | tuple | None = None,
+        missing_value: Any = None,
+    ) -> None:
         """
         Set property vector for main ids.
 
@@ -223,6 +229,7 @@ def set_property(self, key, values: Sequence, ids: Optional[Sequence] = None, mi
             Array of values for the property
         ids : list/np.array, default: None
             List of subset of ids to set the values, default: None
+            if None which is the default all the ids are set or changed
         missing_value : object, default: None
             In case the property is set on a subset of values ("ids" not None),
             it specifies the how the missing values should be filled.
@@ -240,23 +247,26 @@ def set_property(self, key, values: Sequence, ids: Optional[Sequence] = None, mi
         dtype_kind = dtype.kind
 
         if ids is None:
-            assert values.shape[0] == size
+            assert (
+                values.shape[0] == size
+            ), f"Values must have the same size as the main ids: {size} but got array of size {values.shape[0]}"
             self._properties[key] = values
         else:
             ids = np.array(ids)
-            assert np.unique(ids).size == ids.size, "'ids' are not unique!"
+            unique_ids = np.unique(ids)
+            if unique_ids.size != ids.size:
+                non_unique_ids = [id for id in ids if np.count_nonzero(ids == id) > 1]
+                raise ValueError(f"IDs are not unique: {non_unique_ids}")
 
+            # Not clear where this branch is used, perhaps on aggregation of extractors?
             if ids.size < size:
                 if key not in self._properties:
-                    # create the property with nan or empty
-                    shape = (size,) + values.shape[1:]
 
                     if missing_value is None:
                         if dtype_kind not in self.default_missing_property_values.keys():
-                            raise Exception(
-                                "For values dtypes other than float, string, object or unicode, the missing value "
-                                "cannot be automatically inferred. Please specify it with the 'missing_value' "
-                                "argument."
+                            raise ValueError(
+                                f"Can't infer a natural missing value for dtype {dtype_kind}. "
+                                "Please provide it with the `missing_value` argument"
                             )
                         else:
                             missing_value = self.default_missing_property_values[dtype_kind]
@@ -268,15 +278,18 @@ def set_property(self, key, values: Sequence, ids: Optional[Sequence] = None, mi
                             "as the values."
                         )
 
+                    # create the property with nan or empty
+                    shape = (size,) + values.shape[1:]
                     empty_values = np.zeros(shape, dtype=dtype)
                     empty_values[:] = missing_value
                     self._properties[key] = empty_values
                     if ids.size == 0:
                         return
                 else:
-                    assert dtype_kind == self._properties[key].dtype.kind, (
-                        "Mismatch between existing property dtype " "values dtype."
-                    )
+                    existing_property = self._properties[key].dtype
+                    assert (
+                        dtype_kind == existing_property.kind
+                    ), f"Mismatch between existing property dtype {existing_property.kind} and provided values dtype {dtype_kind}."
 
                 indices = self.ids_to_indices(ids)
                 self._properties[key][indices] = values
@@ -285,7 +298,7 @@ def set_property(self, key, values: Sequence, ids: Optional[Sequence] = None, mi
                 self._properties[key] = np.zeros_like(values, dtype=values.dtype)
                 self._properties[key][indices] = values
 
-    def get_property(self, key, ids: Optional[Iterable] = None) -> np.ndarray:
+    def get_property(self, key: str, ids: Optional[Iterable] = None) -> np.ndarray:
         values = self._properties.get(key, None)
         if ids is not None and values is not None:
             inds = self.ids_to_indices(ids)

src/spikeinterface/core/tests/test_base.py

@@ -94,6 +94,41 @@ def test_name_and_repr():
     assert "Hz" in rec_str
 
 
+def test_setting_properties():
+
+    num_channels = 5
+    recording = generate_recording(num_channels=5, durations=[1.0])
+    channel_ids = ["a", "b", "c", "d", "e"]
+    recording = recording.rename_channels(new_channel_ids=channel_ids)
+
+    complete_values = ["value"] * num_channels
+    recording.set_property(key="a_property", values=complete_values)
+
+    property_in_recording = recording.get_property("a_property")
+    expected_array = np.array(complete_values)
+    assert np.array_equal(property_in_recording, expected_array)
+
+    # Set property with missing values
+    incomplete_values = ["value"] * (num_channels - 1)
+    recording.set_property(key="incomplete_property", ids=channel_ids[:-1], values=incomplete_values)
+
+    property_in_recording = recording.get_property("incomplete_property")
+    expected_array = np.array(incomplete_values + [""])  # Spikeinterface defines missing values as empty strings
+    assert np.array_equal(property_in_recording, expected_array)
+
+    # # Passs a missing value
+    # recording.set_property(
+    #     key="missing_property",
+    #     ids=channel_ids[:-1],
+    #     values=incomplete_values,
+    #     missing_value="missing",
+    # )
+
+    # property_in_recording = recording.get_property("missing_property")
+    # expected_array = np.array(incomplete_values + ["missing"])
+    # assert np.array_equal(property_in_recording, expected_array)
+
+
 if __name__ == "__main__":
     test_check_if_memory_serializable()
     test_check_if_serializable()