allow option of different dict per node

Author: TeunHuijben

src/funtracks/user_actions/user_update_nodes_attrs.py

@@ -20,21 +20,30 @@ class UserUpdateNodesAttrs(ActionGroup):
     Args:
         tracks: The tracks to update node attributes for.
         nodes: The node ids to update.
-        attrs: A mapping from attribute name to new attribute values,
-            applied to all nodes.
+        attrs: Either a single dict applied to all nodes, or a list of dicts
+            with one entry per node (must match the length of nodes).
     """
 
     def __init__(
         self,
         tracks: SolutionTracks,
         nodes: list[int],
-        attrs: dict[str, Any],
+        attrs: dict[str, Any] | list[dict[str, Any]],
     ):
         super().__init__(tracks, actions=[])
         self.tracks: SolutionTracks  # Narrow type from base class
-        for node in nodes:
+        if isinstance(attrs, list):
+            if len(attrs) != len(nodes):
+                raise ValueError(
+                    f"attrs list length ({len(attrs)}) must match "
+                    f"nodes length ({len(nodes)})"
+                )
+            per_node_attrs = attrs
+        else:
+            per_node_attrs = [attrs] * len(nodes)
+        for node, node_attrs in zip(nodes, per_node_attrs, strict=True):
             self.actions.append(
-                UserUpdateNodeAttrs(tracks, node, attrs, _top_level=False)
+                UserUpdateNodeAttrs(tracks, node, node_attrs, _top_level=False)
             )
 
         self.tracks.action_history.add_new_action(self)

tests/user_actions/test_user_update_nodes_attrs.py

@@ -45,6 +45,23 @@ def test_undo_redo(self, get_tracks, ndim, with_seg):
         for node in [1, 2]:
             assert tracks.get_node_attr(node, "score") == 0.9
 
+    def test_per_node_attrs(self, get_tracks, ndim, with_seg):
+        """Test bulk update with a different attr dict per node."""
+        tracks = get_tracks(ndim=ndim, with_seg=with_seg, is_solution=True)
+
+        per_node = [{"score": 0.1}, {"score": 0.9}]
+        UserUpdateNodesAttrs(tracks, nodes=[1, 2], attrs=per_node)
+
+        assert tracks.get_node_attr(1, "score") == 0.1
+        assert tracks.get_node_attr(2, "score") == 0.9
+
+    def test_per_node_attrs_length_mismatch_raises(self, get_tracks, ndim, with_seg):
+        """Mismatched list length raises ValueError."""
+        tracks = get_tracks(ndim=ndim, with_seg=with_seg, is_solution=True)
+
+        with pytest.raises(ValueError, match="attrs list length"):
+            UserUpdateNodesAttrs(tracks, nodes=[1, 2], attrs=[{"score": 0.1}])
+
     def test_protected_attr_raises(self, get_tracks, ndim, with_seg):
         """Passing a protected attribute raises ValueError."""
         tracks = get_tracks(ndim=ndim, with_seg=with_seg, is_solution=True)