feat(hooks): add useNodesData, useUpdateNodeData, simplify useHandleConnections

packages/react/src/hooks/useHandleConnections.ts

@@ -4,7 +4,7 @@ import { Connection, HandleType } from '@xyflow/system';
 import { useStore } from './useStore';
 import { useNodeId } from '../contexts/NodeIdContext';
 
-type useHandleConnectionStatusParams = {
+type useHandleConnectionsParams = {
   handleType: HandleType;
   nodeId?: string;
   handleId?: string | null;
@@ -12,6 +12,50 @@ type useHandleConnectionStatusParams = {
   onDisconnect?: (connections: Connection[]) => void;
 };
 
+/**
+ *  Hook to check if a <Handle /> is connected to another <Handle /> and get the connections.
+ *
+ * @public
+ * @param param.handleType - 'source' or 'target'
+ * @param param.handleId - the handle id (this is only needed if the node has multiple handles of the same type)
+ * @param param.nodeId - node id - if not provided, the node id from the NodeIdContext is used
+ * @param param.onConnect - gets called when a connection is established
+ * @param param.onDisconnect - gets called when a connection is removed
+ * @returns an array with connections
+ */
+export function useHandleConnections({
+  handleType,
+  handleId = null,
+  nodeId,
+  onConnect,
+  onDisconnect,
+}: useHandleConnectionsParams): Connection[] {
+  const _nodeId = useNodeId();
+  const prevConnections = useRef<Map<string, Connection> | null>(null);
+  const currentNodeId = nodeId || _nodeId;
+
+  const connections = useStore(
+    (state) => state.connectionLookup.get(`${currentNodeId}-${handleType}-${handleId}`),
+    areConnectionMapsEqual
+  );
+
+  useEffect(() => {
+    // @todo dicuss if onConnect/onDisconnect should be called when the component mounts/unmounts
+    if (prevConnections.current && prevConnections.current !== connections) {
+      const _connections = connections ?? new Map();
+      handleConnectionChange(prevConnections.current, _connections, onDisconnect);
+      handleConnectionChange(_connections, prevConnections.current, onConnect);
+    }
+
+    prevConnections.current = connections ?? new Map();
+  }, [connections, onConnect, onDisconnect]);
+
+  return useMemo(() => Array.from(connections?.values() ?? []), [connections]);
+}
+
+/**
+ * @internal
+ */
 function areConnectionMapsEqual(a?: Map<string, Connection>, b?: Map<string, Connection>) {
   if (!a && !b) {
     return true;
@@ -36,6 +80,7 @@ function areConnectionMapsEqual(a?: Map<string, Connection>, b?: Map<string, Con
 
 /**
  * We call the callback for all connections in a that are not in b
+ *
  * @internal
  */
 function handleConnectionChange(
@@ -59,53 +104,3 @@ function handleConnectionChange(
     cb(diff);
   }
 }
-
-/**
- *  Hook to check if a <Handle /> is connected to another <Handle /> and get the connections.
- *
- * @public
- * @param param.handleType - 'source' or 'target'
- * @param param.handleId - the handle id (this is only needed if the node has multiple handles of the same type)
- * @param param.nodeId - node id - if not provided, the node id from the NodeIdContext is used
- * @param param.onConnect - gets called when a connection is established
- * @param param.onDisconnect - gets called when a connection is removed
- * @returns a `connected` boolean and a connections array
- */
-export function useHandleConnectionStatus({
-  handleType,
-  handleId = null,
-  nodeId,
-  onConnect,
-  onDisconnect,
-}: useHandleConnectionStatusParams): {
-  connected: boolean;
-  connections: Connection[] | null;
-} {
-  const _nodeId = useNodeId();
-  const prevConnections = useRef<Map<string, Connection> | null>(null);
-  const currentNodeId = nodeId || _nodeId;
-
-  const connections = useStore(
-    (state) => state.connectionLookup.get(`${currentNodeId}-${handleType}-${handleId}`),
-    areConnectionMapsEqual
-  );
-
-  useEffect(() => {
-    // @todo dicuss if onConnect/onDisconnect should be called when the component mounts/unmounts
-    if (prevConnections.current && prevConnections.current !== connections) {
-      const _connections = connections ?? new Map();
-      handleConnectionChange(prevConnections.current, _connections, onDisconnect);
-      handleConnectionChange(_connections, prevConnections.current, onConnect);
-    }
-
-    prevConnections.current = connections ?? new Map();
-  }, [connections, onConnect, onDisconnect]);
-
-  return useMemo(
-    () => ({
-      connected: !!connections,
-      connections: Array.from(connections?.values() ?? []),
-    }),
-    [connections]
-  );
-}

packages/react/src/hooks/useNodesData.ts

@@ -0,0 +1,33 @@
+import { useCallback } from 'react';
+import { shallow } from 'zustand/shallow';
+
+import { useStore } from '../hooks/useStore';
+
+export function useNodesData<NodeData = unknown>(nodeId: string): NodeData | null;
+export function useNodesData<NodeData = unknown>(nodeIds: string[]): NodeData[];
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function useNodesData(nodeIds: any): any {
+  const nodesData = useStore(
+    useCallback(
+      (s) => {
+        if (!Array.isArray(nodeIds)) {
+          return s.nodeLookup.get(nodeIds)?.data || null;
+        }
+
+        return nodeIds.reduce((res, id) => {
+          const node = s.nodeLookup.get(id);
+
+          if (node) {
+            res.push(node.data);
+          }
+
+          return res;
+        }, []);
+      },
+      [nodeIds]
+    ),
+    shallow
+  );
+
+  return nodesData;
+}

packages/react/src/hooks/useUpdateNodeData.ts

@@ -0,0 +1,13 @@
+import { useCallback } from 'react';
+
+import useReactFlow from './useReactFlow';
+
+export function useUpdateNodeData() {
+  const { setNodes } = useReactFlow();
+
+  const updateNodeData = useCallback((id: string, data: unknown) => {
+    setNodes((prevNodes) => prevNodes.map((node) => (node.id === id ? { ...node, data } : node)));
+  }, []);
+
+  return updateNodeData;
+}

packages/react/src/index.ts

@@ -1,5 +1,5 @@
 export { default as ReactFlow } from './container/ReactFlow';
-export { default as Handle } from './components/Handle';
+export { default as Handle, type HandleComponentProps } from './components/Handle';
 export { default as EdgeText } from './components/Edges/EdgeText';
 export { StraightEdge } from './components/Edges/StraightEdge';
 export { StepEdge } from './components/Edges/StepEdge';
@@ -22,7 +22,9 @@ export { useStore, useStoreApi } from './hooks/useStore';
 export { default as useOnViewportChange, type UseOnViewportChangeOptions } from './hooks/useOnViewportChange';
 export { default as useOnSelectionChange, type UseOnSelectionChangeOptions } from './hooks/useOnSelectionChange';
 export { default as useNodesInitialized, type UseNodesInitializedOptions } from './hooks/useNodesInitialized';
-export { useHandleConnectionStatus } from './hooks/useHandleConnectionStatus';
+export { useHandleConnections } from './hooks/useHandleConnections';
+export { useUpdateNodeData } from './hooks/useUpdateNodeData';
+export { useNodesData } from './hooks/useNodesData';
 export { useNodeId } from './contexts/NodeIdContext';
 
 export { applyNodeChanges, applyEdgeChanges, handleParentExpand } from './utils/changes';