From 3124d4673259bf1b326e4b05c47f8abee2995b2a Mon Sep 17 00:00:00 2001
From: moklick <info@moritzklack.com>
Date: Fri, 28 May 2021 14:26:03 +0200
Subject: [PATCH] refactor(edge-updater): pass handle type #1212

---
 example/src/UpdatableEdge/index.tsx | 3 +++
 src/components/Edges/wrapEdge.tsx   | 3 ++-
 src/components/Handle/handler.ts    | 7 ++++---
 src/components/Handle/index.tsx     | 5 +++--
 4 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/example/src/UpdatableEdge/index.tsx b/example/src/UpdatableEdge/index.tsx
index 1b0920e5..b6a40547 100644
--- a/example/src/UpdatableEdge/index.tsx
+++ b/example/src/UpdatableEdge/index.tsx
@@ -7,6 +7,7 @@ import ReactFlow, {
   OnLoadParams,
   Connection,
   Edge,
+  removeElements,
 } from 'react-flow-renderer';
 
 const initialElements: Elements = [
@@ -56,6 +57,7 @@ const UpdatableEdge = () => {
   const onEdgeUpdate = (oldEdge: Edge, newConnection: Connection) =>
     setElements((els) => updateEdge(oldEdge, newConnection, els));
   const onConnect = (params: Connection | Edge) => setElements((els) => addEdge(params, els));
+  const onElementsRemove = (elementsToRemove: Elements) => setElements((els) => removeElements(elementsToRemove, els));
 
   return (
     <ReactFlow
@@ -65,6 +67,7 @@ const UpdatableEdge = () => {
       onEdgeUpdate={onEdgeUpdate}
       onConnect={onConnect}
       onEdgeUpdateStart={onEdgeUpdateStart}
+      onElementsRemove={onElementsRemove}
     >
       <Controls />
     </ReactFlow>
diff --git a/src/components/Edges/wrapEdge.tsx b/src/components/Edges/wrapEdge.tsx
index db519135..7d5274a6 100644
--- a/src/components/Edges/wrapEdge.tsx
+++ b/src/components/Edges/wrapEdge.tsx
@@ -150,7 +150,8 @@ export default (EdgeComponent: ComponentType<EdgeProps>) => {
           onConnectEdge,
           isTarget,
           isValidConnection,
-          connectionMode
+          connectionMode,
+          isSourceHandle ? 'target' : 'source'
         );
       },
       [id, source, target, type, sourceHandleId, targetHandleId, setConnectionNodeId, setPosition, edgeElement]
diff --git a/src/components/Handle/handler.ts b/src/components/Handle/handler.ts
index 7f5307e7..0a852c36 100644
--- a/src/components/Handle/handler.ts
+++ b/src/components/Handle/handler.ts
@@ -12,6 +12,7 @@ import {
   ConnectionMode,
   SetConnectionId,
   Connection,
+  HandleType,
 } from '../../types';
 
 type ValidConnectionFunc = (connection: Connection) => boolean;
@@ -96,6 +97,7 @@ export function onMouseDown(
   isTarget: boolean,
   isValidConnection: ValidConnectionFunc,
   connectionMode: ConnectionMode,
+  elementEdgeUpdaterType?: HandleType | null,
   onConnectStart?: OnConnectStartFunc,
   onConnectStop?: OnConnectStopFunc,
   onConnectEnd?: OnConnectEndFunc
@@ -111,13 +113,12 @@ export function onMouseDown(
   const elementBelow = doc.elementFromPoint(event.clientX, event.clientY);
   const elementBelowIsTarget = elementBelow?.classList.contains('target');
   const elementBelowIsSource = elementBelow?.classList.contains('source');
-  const elementBelowIsUpdater = elementBelow?.classList.contains('react-flow__edgeupdater');
 
-  if (!reactFlowNode || (!elementBelowIsTarget && !elementBelowIsSource && !elementBelowIsUpdater)) {
+  if (!reactFlowNode || (!elementBelowIsTarget && !elementBelowIsSource && !elementEdgeUpdaterType)) {
     return;
   }
 
-  const handleType = elementBelowIsTarget ? 'target' : 'source';
+  const handleType = elementEdgeUpdaterType ? elementEdgeUpdaterType : elementBelowIsTarget ? 'target' : 'source';
   const containerBounds = reactFlowNode.getBoundingClientRect();
   let recentHoveredHandle: Element;
 
diff --git a/src/components/Handle/index.tsx b/src/components/Handle/index.tsx
index 0b6da001..3f2aca8d 100644
--- a/src/components/Handle/index.tsx
+++ b/src/components/Handle/index.tsx
@@ -45,12 +45,13 @@ const Handle: FC<HandleProps & Omit<HTMLAttributes<HTMLDivElement>, 'id'>> = ({
         event,
         handleId,
         nodeId,
-        (setConnectionNodeId as unknown) as SetSourceIdFunc,
-        (setPosition as unknown) as SetPosition,
+        setConnectionNodeId as unknown as SetSourceIdFunc,
+        setPosition as unknown as SetPosition,
         onConnectExtended,
         isTarget,
         isValidConnection,
         connectionMode,
+        null,
         onConnectStart,
         onConnectStop,
         onConnectEnd