Merge pull request #4554 from xyflow/use-connection-selector

added optional selector to useConnection hook

.changeset/tender-jobs-kneel.md

@@ -0,0 +1,5 @@
+---
+'@xyflow/react': minor
+---
+
+Added optional selector for useConnection hook

packages/react/src/hooks/useConnection.ts

@@ -4,17 +4,34 @@ import { ConnectionState, pointToRendererPoint } from '@xyflow/system';
 import { useStore } from './useStore';
 import type { InternalNode, Node, ReactFlowStore } from '../types';
 
-const selector = (s: ReactFlowStore) => {
+function storeSelector(s: ReactFlowStore) {
   return s.connection.inProgress
     ? { ...s.connection, to: pointToRendererPoint(s.connection.to, s.transform) }
     : { ...s.connection };
-};
+}
+
+function getSelector<NodeType extends Node = Node, SelectorReturn = ConnectionState<InternalNode<NodeType>>>(
+  connectionSelector?: (connection: ConnectionState<InternalNode<NodeType>>) => SelectorReturn
+): (s: ReactFlowStore) => SelectorReturn | ConnectionState<InternalNode> {
+  if (connectionSelector) {
+    const combinedSelector = (s: ReactFlowStore) => {
+      const connection = storeSelector(s) as ConnectionState<InternalNode<NodeType>>;
+      return connectionSelector(connection);
+    };
+    return combinedSelector;
+  }
+
+  return storeSelector;
+}
 /**
  * Hook for accessing the connection state.
  *
  * @public
  * @returns ConnectionState
  */
-export function useConnection<NodeType extends Node = Node>(): ConnectionState<InternalNode<NodeType>> {
-  return useStore(selector, shallow) as ConnectionState<InternalNode<NodeType>>;
+export function useConnection<NodeType extends Node = Node, SelectorReturn = ConnectionState<InternalNode<NodeType>>>(
+  connectionSelector?: (connection: ConnectionState<InternalNode<NodeType>>) => SelectorReturn
+): SelectorReturn {
+  const combinedSelector = getSelector<NodeType, SelectorReturn>(connectionSelector);
+  return useStore(combinedSelector, shallow) as SelectorReturn;
 }