feat(useReactFlow): add updateEdge and updateEdgeData (#4351)

* feat(useReactFlow): add updateEdge and updateEdgeData

* chore(examples): add updateEdge call

examples/react/src/examples/MultiSetNodes/index.tsx

@@ -28,10 +28,15 @@ for (let i = 0; i < 100; i++) {
   });
 }
 
-const initEdges: Edge[] = [];
+const initEdges: Edge[] = initNodes.reduce<Edge[]>((res, node, index) => {
+  if (index > 0) {
+    res.push({ id: `${index - 1}-${index}`, source: (index - 1).toString(), target: index.toString() });
+  }
+  return res;
+}, []);
 
 const CustomNodeFlow = () => {
-  const { setNodes, updateNodeData } = useReactFlow();
+  const { setNodes, updateNodeData, updateEdge } = useReactFlow();
   const [nodes, , onNodesChange] = useNodesState(initNodes);
   const [edges, setEdges, onEdgesChange] = useEdgesState(initEdges);
 
@@ -55,6 +60,10 @@ const CustomNodeFlow = () => {
     nodes.forEach((node) => updateNodeData(node.id, { label: 'node update' }));
   };
 
+  const multiUpdateEdges = () => {
+    edges.forEach((edge) => updateEdge(edge.id, { label: 'edge update' }));
+  };
+
   return (
     <ReactFlow
       nodes={nodes}
@@ -70,6 +79,7 @@ const CustomNodeFlow = () => {
       <Panel>
         <button onClick={multiSetNodes}>set nodes</button>
         <button onClick={multiUpdateNodes}>update nodes</button>
+        <button onClick={multiUpdateEdges}>update edges</button>
       </Panel>
     </ReactFlow>
   );

packages/react/src/hooks/useReactFlow.ts

@@ -13,7 +13,7 @@ import {
 import useViewportHelper from './useViewportHelper';
 import { useStore, useStoreApi } from './useStore';
 import { useBatchContext } from '../components/BatchProvider';
-import { elementToRemoveChange, isNode } from '../utils';
+import { elementToRemoveChange, isEdge, isNode } from '../utils';
 import type { ReactFlowInstance, Node, Edge, InternalNode, ReactFlowState, GeneralHelpers } from '../types';
 
 const selector = (s: ReactFlowState) => !!s.panZoom;
@@ -41,6 +41,10 @@ export function useReactFlow<NodeType extends Node = Node, EdgeType extends Edge
       batchContext.nodeQueue.push(payload as NodeType[]);
     };
 
+    const setEdges: GeneralHelpers<NodeType, EdgeType>['setEdges'] = (payload) => {
+      batchContext.edgeQueue.push(payload as EdgeType[]);
+    };
+
     const getNodeRect = (node: NodeType | { id: string }): Rect | null => {
       const { nodeLookup, nodeOrigin } = store.getState();
 
@@ -77,6 +81,23 @@ export function useReactFlow<NodeType extends Node = Node, EdgeType extends Edge
       );
     };
 
+    const updateEdge: GeneralHelpers<NodeType, EdgeType>['updateEdge'] = (
+      id,
+      edgeUpdate,
+      options = { replace: false }
+    ) => {
+      setEdges((prevEdges) =>
+        prevEdges.map((edge) => {
+          if (edge.id === id) {
+            const nextEdge = typeof edgeUpdate === 'function' ? edgeUpdate(edge as EdgeType) : edgeUpdate;
+            return options.replace && isEdge(nextEdge) ? (nextEdge as EdgeType) : { ...edge, ...nextEdge };
+          }
+
+          return edge;
+        })
+      );
+    };
+
     return {
       getNodes: () => store.getState().nodes.map((n) => ({ ...n })) as NodeType[],
       getNode: (id) => getInternalNode(id)?.internals.userNode as NodeType,
@@ -87,9 +108,7 @@ export function useReactFlow<NodeType extends Node = Node, EdgeType extends Edge
       },
       getEdge: (id) => store.getState().edgeLookup.get(id) as EdgeType,
       setNodes,
-      setEdges: (payload) => {
-        batchContext.edgeQueue.push(payload as EdgeType[]);
-      },
+      setEdges,
       addNodes: (payload) => {
         const newNodes = Array.isArray(payload) ? payload : [payload];
         batchContext.nodeQueue.push((nodes) => [...nodes, ...newNodes]);
@@ -200,6 +219,17 @@ export function useReactFlow<NodeType extends Node = Node, EdgeType extends Edge
           options
         );
       },
+      updateEdge,
+      updateEdgeData: (id, dataUpdate, options = { replace: false }) => {
+        updateEdge(
+          id,
+          (edge) => {
+            const nextData = typeof dataUpdate === 'function' ? dataUpdate(edge) : dataUpdate;
+            return options.replace ? { ...edge, data: nextData } : { ...edge, data: { ...edge.data, ...nextData } };
+          },
+          options
+        );
+      },
     };
   }, []);
 

packages/react/src/types/instance.ts

@@ -143,6 +143,36 @@ export type GeneralHelpers<NodeType extends Node = Node, EdgeType extends Edge =
     dataUpdate: Partial<NodeType['data']> | ((node: NodeType) => Partial<NodeType['data']>),
     options?: { replace: boolean }
   ) => void;
+  /**
+   * Updates an edge.
+   *
+   * @param id - id of the edge to update
+   * @param edgeUpdate - the edge update as an object or a function that receives the current edge and returns the edge update
+   * @param options.replace - if true, the edge is replaced with the edge update, otherwise the changes get merged
+   *
+   * @example
+   * updateEdge('edge-1', (edge) => ({ label: 'A new label' }));
+   */
+  updateEdge: (
+    id: string,
+    edgeUpdate: Partial<EdgeType> | ((edge: EdgeType) => Partial<EdgeType>),
+    options?: { replace: boolean }
+  ) => void;
+  /**
+   * Updates the data attribute of a edge.
+   *
+   * @param id - id of the edge to update
+   * @param dataUpdate - the data update as an object or a function that receives the current data and returns the data update
+   * @param options.replace - if true, the data is replaced with the data update, otherwise the changes get merged
+   *
+   * @example
+   * updateEdgeData('edge-1', { label: 'A new label' });
+   */
+  updateEdgeData: (
+    id: string,
+    dataUpdate: Partial<EdgeType['data']> | ((edge: EdgeType) => Partial<EdgeType['data']>),
+    options?: { replace: boolean }
+  ) => void;
 };
 
 export type ReactFlowInstance<NodeType extends Node = Node, EdgeType extends Edge = Edge> = GeneralHelpers<