Merge pull request #2877 from wbkd/feat/isValidConnection

feat(reactflow): add isValidConnection prop

.changeset/thick-garlics-tap.md

@@ -0,0 +1,5 @@
+---
+'@reactflow/core': minor
+---
+
+add isValidConnection prop for ReactFlow component

examples/vite-app/src/examples/Validation/index.tsx

@@ -12,6 +12,8 @@ import ReactFlow, {
   OnConnectStart,
   OnConnectEnd,
   OnConnect,
+  updateEdge,
+  Edge,
 } from 'reactflow';
 
 import styles from './validation.module.css';
@@ -28,15 +30,15 @@ const isValidConnection = (connection: Connection) => connection.target === 'B';
 const CustomInput: FC<NodeProps> = () => (
   <>
     <div>Only connectable with B</div>
-    <Handle type="source" position={Position.Right} isValidConnection={isValidConnection} />
+    <Handle type="source" position={Position.Right} />
   </>
 );
 
 const CustomNode: FC<NodeProps> = ({ id }) => (
   <>
-    <Handle type="target" position={Position.Left} isValidConnection={isValidConnection} />
+    <Handle type="target" position={Position.Left} />
     <div>{id}</div>
-    <Handle type="source" position={Position.Right} isValidConnection={isValidConnection} />
+    <Handle type="source" position={Position.Right} />
   </>
 );
 
@@ -74,6 +76,11 @@ const ValidationFlow = () => {
     [value]
   );
 
+  const onEdgeUpdate = useCallback(
+    (oldEdge: Edge, newConnection: Connection) => setEdges((els) => updateEdge(oldEdge, newConnection, els)),
+    [setEdges]
+  );
+
   return (
     <ReactFlow
       nodes={nodes}
@@ -86,6 +93,8 @@ const ValidationFlow = () => {
       nodeTypes={nodeTypes}
       onConnectStart={onConnectStart}
       onConnectEnd={onConnectEnd}
+      onEdgeUpdate={onEdgeUpdate}
+      isValidConnection={isValidConnection}
       fitView
     />
   );

packages/core/src/components/Edges/wrapEdge.tsx

@@ -11,6 +11,8 @@ import { getMouseHandler } from './utils';
 import { elementSelectionKeys } from '../../utils';
 import type { EdgeProps, WrapEdgeProps, Connection } from '../../types';
 
+const alwaysValidConnection = () => true;
+
 export default (EdgeComponent: ComponentType<EdgeProps>) => {
   const EdgeWrapper = ({
     id,
@@ -94,12 +96,14 @@ export default (EdgeComponent: ComponentType<EdgeProps>) => {
         return;
       }
 
+      const { edges, isValidConnection: isValidConnectionStore } = store.getState();
       const nodeId = isSourceHandle ? target : source;
       const handleId = (isSourceHandle ? targetHandleId : sourceHandleId) || null;
       const handleType = isSourceHandle ? 'target' : 'source';
-      const isValidConnection = () => true;
+      const isValidConnection = isValidConnectionStore || alwaysValidConnection;
+
       const isTarget = isSourceHandle;
-      const edge = store.getState().edges.find((e) => e.id === id)!;
+      const edge = edges.find((e) => e.id === id)!;
 
       setUpdating(true);
       onEdgeUpdateStart?.(event, edge, handleType);

packages/core/src/components/Handle/index.tsx

@@ -27,7 +27,7 @@ const Handle = forwardRef<HTMLDivElement, HandleComponentProps>(
     {
       type = 'source',
       position = Position.Top,
-      isValidConnection = alwaysValid,
+      isValidConnection,
       isConnectable = true,
       id,
       onConnect,
@@ -61,8 +61,8 @@ const Handle = forwardRef<HTMLDivElement, HandleComponentProps>(
         ...params,
       };
       if (hasDefaultEdges) {
-        const { edges } = store.getState();
-        store.setState({ edges: addEdge(edgeParams, edges) });
+        const { edges, setEdges } = store.getState();
+        setEdges(addEdge(edgeParams, edges));
       }
 
       onConnectAction?.(edgeParams);
@@ -81,7 +81,7 @@ const Handle = forwardRef<HTMLDivElement, HandleComponentProps>(
           isTarget,
           getState: store.getState,
           setState: store.setState,
-          isValidConnection,
+          isValidConnection: isValidConnection || store.getState().isValidConnection || alwaysValid,
         });
       }
 
@@ -93,7 +93,12 @@ const Handle = forwardRef<HTMLDivElement, HandleComponentProps>(
     };
 
     const onClick = (event: ReactMouseEvent) => {
-      const { onClickConnectStart, onClickConnectEnd, connectionMode } = store.getState();
+      const {
+        onClickConnectStart,
+        onClickConnectEnd,
+        connectionMode,
+        isValidConnection: isValidConnectionStore,
+      } = store.getState();
       if (!connectionStartHandle) {
         onClickConnectStart?.(event, { nodeId, handleId, handleType: type });
         store.setState({ connectionStartHandle: { nodeId, type, handleId } });
@@ -101,6 +106,7 @@ const Handle = forwardRef<HTMLDivElement, HandleComponentProps>(
       }
 
       const doc = getHostForElement(event.target as HTMLElement);
+      const isValidConnectionHandler = isValidConnection || isValidConnectionStore || alwaysValid;
       const { connection, isValid } = isValidHandle(
         event,
         {
@@ -112,7 +118,7 @@ const Handle = forwardRef<HTMLDivElement, HandleComponentProps>(
         connectionStartHandle.nodeId,
         connectionStartHandle.handleId || null,
         connectionStartHandle.type,
-        isValidConnection,
+        isValidConnectionHandler,
         doc
       );
 

packages/core/src/components/StoreUpdater/index.tsx

@@ -49,6 +49,7 @@ type StoreUpdaterProps = Pick<
   | 'autoPanOnNodeDrag'
   | 'onError'
   | 'connectionRadius'
+  | 'isValidConnection'
 > & { rfId: string };
 
 const selector = (s: ReactFlowState) => ({
@@ -127,6 +128,7 @@ const StoreUpdater = ({
   autoPanOnNodeDrag,
   onError,
   connectionRadius,
+  isValidConnection,
 }: StoreUpdaterProps) => {
   const {
     setNodes,
@@ -184,6 +186,7 @@ const StoreUpdater = ({
   useDirectStoreUpdater('autoPanOnNodeDrag', autoPanOnNodeDrag, store.setState);
   useDirectStoreUpdater('onError', onError, store.setState);
   useDirectStoreUpdater('connectionRadius', connectionRadius, store.setState);
+  useDirectStoreUpdater('isValidConnection', isValidConnection, store.setState);
 
   useStoreUpdater<Node[]>(nodes, setNodes);
   useStoreUpdater<Edge[]>(edges, setEdges);

packages/core/src/container/ReactFlow/index.tsx

@@ -163,6 +163,7 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
       autoPanOnConnect = true,
       autoPanOnNodeDrag = true,
       connectionRadius = 20,
+      isValidConnection,
       onError,
       style,
       id,
@@ -295,6 +296,7 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
             autoPanOnNodeDrag={autoPanOnNodeDrag}
             onError={onError}
             connectionRadius={connectionRadius}
+            isValidConnection={isValidConnection}
           />
           <SelectionListener onSelectionChange={onSelectionChange} />
           {children}

packages/core/src/store/index.ts

@@ -33,8 +33,8 @@ const createRFStore = () =>
       return Array.from(get().nodeInternals.values());
     },
     setEdges: (edges: Edge[]) => {
-      const { defaultEdgeOptions = {} } = get();
-      set({ edges: edges.map((e) => ({ ...defaultEdgeOptions, ...e })) });
+      const { defaultEdgeOptions = null } = get();
+      set({ edges: defaultEdgeOptions ? edges.map((e) => ({ ...defaultEdgeOptions, ...e })) : edges });
     },
     setDefaultNodesAndEdges: (nodes?: Node[], edges?: Edge[]) => {
       const hasDefaultNodes = typeof nodes !== 'undefined';

packages/core/src/store/initialState.ts

@@ -62,6 +62,7 @@ const initialState: ReactFlowStore = {
   autoPanOnNodeDrag: true,
   connectionRadius: 20,
   onError: devWarn,
+  isValidConnection: undefined,
 };
 
 export default initialState;

packages/core/src/types/component-props.ts

@@ -38,6 +38,7 @@ import type {
   SelectionMode,
   OnError,
 } from '.';
+import { ValidConnectionFunc } from '../components/Handle/utils';
 
 export type ReactFlowProps = HTMLAttributes<HTMLDivElement> & {
   nodes?: Node[];
@@ -144,6 +145,7 @@ export type ReactFlowProps = HTMLAttributes<HTMLDivElement> & {
   autoPanOnConnect?: boolean;
   connectionRadius?: number;
   onError?: OnError;
+  isValidConnection?: ValidConnectionFunc;
 };
 
 export type ReactFlowRefType = HTMLDivElement;

packages/core/src/types/general.ts

@@ -61,6 +61,8 @@ export interface Connection {
   targetHandle: string | null;
 }
 
+export type IsValidConnection = (edge: Edge | Connection) => boolean;
+
 export type ConnectionStatus = 'valid' | 'invalid';
 
 export enum ConnectionMode {
@@ -223,6 +225,8 @@ export type ReactFlowStore = {
   autoPanOnConnect: boolean;
   autoPanOnNodeDrag: boolean;
   connectionRadius: number;
+
+  isValidConnection?: IsValidConnection;
 };
 
 export type ReactFlowActions = {