refactor(edges): use EdgeWrapper instead of wrapEdge

packages/react/src/components/EdgeWrapper/index.tsx

@@ -0,0 +1,307 @@
+import { memo, useState, useMemo, useRef, type KeyboardEvent, useCallback } from 'react';
+import cc from 'classcat';
+import { shallow } from 'zustand/shallow';
+import {
+  getMarkerId,
+  elementSelectionKeys,
+  XYHandle,
+  type Connection,
+  getEdgePosition,
+  errorMessages,
+} from '@xyflow/system';
+
+import { useStoreApi, useStore } from '../../hooks/useStore';
+import { ARIA_EDGE_DESC_KEY } from '../A11yDescriptions';
+import { EdgeAnchor } from '../Edges/EdgeAnchor';
+import { getMouseHandler } from '../Edges/utils';
+import type { EdgeWrapperProps } from '../../types';
+import { builtinEdgeTypes } from './utils';
+
+function EdgeWrapper({
+  id,
+  className,
+  type,
+  data,
+  onClick,
+  onEdgeDoubleClick,
+  selected,
+  animated,
+  label,
+  labelStyle,
+  labelShowBg,
+  labelBgStyle,
+  labelBgPadding,
+  labelBgBorderRadius,
+  style,
+  source,
+  target,
+  isSelectable,
+  hidden,
+  sourceHandleId,
+  targetHandleId,
+  onContextMenu,
+  onMouseEnter,
+  onMouseMove,
+  onMouseLeave,
+  edgeUpdaterRadius,
+  onEdgeUpdate,
+  onEdgeUpdateStart,
+  onEdgeUpdateEnd,
+  markerEnd,
+  markerStart,
+  rfId,
+  ariaLabel,
+  isFocusable,
+  isUpdatable,
+  pathOptions,
+  interactionWidth,
+  edgeTypes,
+  onError,
+}: EdgeWrapperProps): JSX.Element | null {
+  let edgeType = type || 'default';
+  let EdgeComponent = edgeTypes?.[edgeType] || builtinEdgeTypes[edgeType];
+
+  if (EdgeComponent === undefined) {
+    onError?.('011', errorMessages['error011'](edgeType));
+    edgeType = 'default';
+    EdgeComponent = builtinEdgeTypes.default;
+  }
+
+  const edgeRef = useRef<SVGGElement>(null);
+  const [updateHover, setUpdateHover] = useState<boolean>(false);
+  const [updating, setUpdating] = useState<boolean>(false);
+  const store = useStoreApi();
+  const edgePosition = useStore(
+    useCallback(
+      (state) => {
+        const sourceNode = state.nodeLookup.get(source);
+        const targetNode = state.nodeLookup.get(target);
+
+        if (!sourceNode || !targetNode) {
+          return null;
+        }
+
+        return getEdgePosition({
+          id,
+          sourceNode,
+          targetNode,
+          sourceHandle: sourceHandleId || null,
+          targetHandle: targetHandleId || null,
+          connectionMode: state.connectionMode,
+          onError: state.onError,
+        });
+      },
+      [source, target]
+    ),
+    shallow
+  );
+
+  const markerStartUrl = useMemo(() => `url(#${getMarkerId(markerStart, rfId)})`, [markerStart, rfId]);
+  const markerEndUrl = useMemo(() => `url(#${getMarkerId(markerEnd, rfId)})`, [markerEnd, rfId]);
+
+  if (hidden || !edgePosition) {
+    return null;
+  }
+
+  const onEdgeClick = (event: React.MouseEvent<SVGGElement, MouseEvent>): void => {
+    const { edges, addSelectedEdges, unselectNodesAndEdges, multiSelectionActive } = store.getState();
+    const edge = edges.find((e) => e.id === id);
+
+    if (!edge) {
+      return;
+    }
+
+    if (isSelectable) {
+      store.setState({ nodesSelectionActive: false });
+
+      if (edge.selected && multiSelectionActive) {
+        unselectNodesAndEdges({ nodes: [], edges: [edge] });
+        edgeRef.current?.blur();
+      } else {
+        addSelectedEdges([id]);
+      }
+    }
+
+    if (onClick) {
+      onClick(event, edge);
+    }
+  };
+
+  const onEdgeDoubleClickHandler = getMouseHandler(id, store.getState, onEdgeDoubleClick);
+  const onEdgeContextMenu = getMouseHandler(id, store.getState, onContextMenu);
+  const onEdgeMouseEnter = getMouseHandler(id, store.getState, onMouseEnter);
+  const onEdgeMouseMove = getMouseHandler(id, store.getState, onMouseMove);
+  const onEdgeMouseLeave = getMouseHandler(id, store.getState, onMouseLeave);
+
+  const handleEdgeUpdater = (event: React.MouseEvent<SVGGElement, MouseEvent>, isSourceHandle: boolean) => {
+    // avoid triggering edge updater if mouse btn is not left
+    if (event.button !== 0) {
+      return;
+    }
+
+    const {
+      autoPanOnConnect,
+      domNode,
+      edges,
+      isValidConnection,
+      connectionMode,
+      connectionRadius,
+      lib,
+      onConnectStart,
+      onConnectEnd,
+      cancelConnection,
+      nodes,
+      panBy,
+      updateConnection,
+    } = store.getState();
+    const nodeId = isSourceHandle ? target : source;
+    const handleId = (isSourceHandle ? targetHandleId : sourceHandleId) || null;
+    const handleType = isSourceHandle ? 'target' : 'source';
+
+    const isTarget = isSourceHandle;
+    const edge = edges.find((e) => e.id === id)!;
+
+    setUpdating(true);
+    onEdgeUpdateStart?.(event, edge, handleType);
+
+    const _onEdgeUpdateEnd = (evt: MouseEvent | TouchEvent) => {
+      setUpdating(false);
+      onEdgeUpdateEnd?.(evt, edge, handleType);
+    };
+
+    const onConnectEdge = (connection: Connection) => onEdgeUpdate?.(edge, connection);
+
+    XYHandle.onPointerDown(event.nativeEvent, {
+      autoPanOnConnect,
+      connectionMode,
+      connectionRadius,
+      domNode,
+      handleId,
+      nodeId,
+      nodes,
+      isTarget,
+      edgeUpdaterType: handleType,
+      lib,
+      cancelConnection,
+      panBy,
+      isValidConnection,
+      onConnect: onConnectEdge,
+      onConnectStart,
+      onConnectEnd,
+      onEdgeUpdateEnd: _onEdgeUpdateEnd,
+      updateConnection,
+      getTransform: () => store.getState().transform,
+    });
+  };
+
+  const onEdgeUpdaterSourceMouseDown = (event: React.MouseEvent<SVGGElement, MouseEvent>): void =>
+    handleEdgeUpdater(event, true);
+  const onEdgeUpdaterTargetMouseDown = (event: React.MouseEvent<SVGGElement, MouseEvent>): void =>
+    handleEdgeUpdater(event, false);
+
+  const onEdgeUpdaterMouseEnter = () => setUpdateHover(true);
+  const onEdgeUpdaterMouseOut = () => setUpdateHover(false);
+
+  const inactive = !isSelectable && !onClick;
+
+  const onKeyDown = (event: KeyboardEvent) => {
+    if (elementSelectionKeys.includes(event.key) && isSelectable) {
+      const { unselectNodesAndEdges, addSelectedEdges, edges } = store.getState();
+      const unselect = event.key === 'Escape';
+
+      if (unselect) {
+        edgeRef.current?.blur();
+        unselectNodesAndEdges({ edges: [edges.find((e) => e.id === id)!] });
+      } else {
+        addSelectedEdges([id]);
+      }
+    }
+  };
+
+  return (
+    <g
+      className={cc([
+        'react-flow__edge',
+        `react-flow__edge-${type}`,
+        className,
+        { selected, animated, inactive, updating: updateHover },
+      ])}
+      onClick={onEdgeClick}
+      onDoubleClick={onEdgeDoubleClickHandler}
+      onContextMenu={onEdgeContextMenu}
+      onMouseEnter={onEdgeMouseEnter}
+      onMouseMove={onEdgeMouseMove}
+      onMouseLeave={onEdgeMouseLeave}
+      onKeyDown={isFocusable ? onKeyDown : undefined}
+      tabIndex={isFocusable ? 0 : undefined}
+      role={isFocusable ? 'button' : 'img'}
+      data-id={id}
+      data-testid={`rf__edge-${id}`}
+      aria-label={ariaLabel === null ? undefined : ariaLabel ? ariaLabel : `Edge from ${source} to ${target}`}
+      aria-describedby={isFocusable ? `${ARIA_EDGE_DESC_KEY}-${rfId}` : undefined}
+      ref={edgeRef}
+    >
+      {!updating && (
+        <EdgeComponent
+          id={id}
+          source={source}
+          target={target}
+          selected={selected}
+          animated={animated}
+          label={label}
+          labelStyle={labelStyle}
+          labelShowBg={labelShowBg}
+          labelBgStyle={labelBgStyle}
+          labelBgPadding={labelBgPadding}
+          labelBgBorderRadius={labelBgBorderRadius}
+          data={data}
+          style={style}
+          sourceX={edgePosition.sourceX}
+          sourceY={edgePosition.sourceY}
+          targetX={edgePosition.targetX}
+          targetY={edgePosition.targetY}
+          sourcePosition={edgePosition.sourcePosition}
+          targetPosition={edgePosition.targetPosition}
+          sourceHandleId={sourceHandleId}
+          targetHandleId={targetHandleId}
+          markerStart={markerStartUrl}
+          markerEnd={markerEndUrl}
+          pathOptions={pathOptions}
+          interactionWidth={interactionWidth}
+        />
+      )}
+      {isUpdatable && (
+        <>
+          {(isUpdatable === 'source' || isUpdatable === true) && (
+            <EdgeAnchor
+              position={edgePosition.sourcePosition}
+              centerX={edgePosition.sourceX}
+              centerY={edgePosition.sourceY}
+              radius={edgeUpdaterRadius}
+              onMouseDown={onEdgeUpdaterSourceMouseDown}
+              onMouseEnter={onEdgeUpdaterMouseEnter}
+              onMouseOut={onEdgeUpdaterMouseOut}
+              type="source"
+            />
+          )}
+          {(isUpdatable === 'target' || isUpdatable === true) && (
+            <EdgeAnchor
+              position={edgePosition.targetPosition}
+              centerX={edgePosition.targetX}
+              centerY={edgePosition.targetY}
+              radius={edgeUpdaterRadius}
+              onMouseDown={onEdgeUpdaterTargetMouseDown}
+              onMouseEnter={onEdgeUpdaterMouseEnter}
+              onMouseOut={onEdgeUpdaterMouseOut}
+              type="target"
+            />
+          )}
+        </>
+      )}
+    </g>
+  );
+}
+
+EdgeWrapper.displayName = 'EdgeWrapper';
+
+export default memo(EdgeWrapper);

packages/react/src/components/EdgeWrapper/utils.ts

@@ -0,0 +1,17 @@
+import type { ComponentType } from 'react';
+import type { EdgeProps, EdgeTypes } from '../../types';
+import {
+  BezierEdgeInternal,
+  StraightEdgeInternal,
+  StepEdgeInternal,
+  SmoothStepEdgeInternal,
+  SimpleBezierEdgeInternal,
+} from '../Edges';
+
+export const builtinEdgeTypes: EdgeTypes = {
+  default: BezierEdgeInternal as ComponentType<EdgeProps>,
+  straight: StraightEdgeInternal as ComponentType<EdgeProps>,
+  step: StepEdgeInternal as ComponentType<EdgeProps>,
+  smoothstep: SmoothStepEdgeInternal as ComponentType<EdgeProps>,
+  simplebezier: SimpleBezierEdgeInternal as ComponentType<EdgeProps>,
+};

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

@@ -1,290 +0,0 @@
-import { memo, useState, useMemo, useRef, type ComponentType, type KeyboardEvent, useCallback } from 'react';
-import cc from 'classcat';
-import { shallow } from 'zustand/shallow';
-import { getMarkerId, elementSelectionKeys, XYHandle, type Connection, getEdgePosition } from '@xyflow/system';
-
-import { useStoreApi, useStore } from '../../hooks/useStore';
-import { ARIA_EDGE_DESC_KEY } from '../A11yDescriptions';
-import { EdgeAnchor } from './EdgeAnchor';
-import { getMouseHandler } from './utils';
-import type { EdgeProps, WrapEdgeProps } from '../../types';
-
-export default (EdgeComponent: ComponentType<EdgeProps>) => {
-  const EdgeWrapper = ({
-    id,
-    className,
-    type,
-    data,
-    onClick,
-    onEdgeDoubleClick,
-    selected,
-    animated,
-    label,
-    labelStyle,
-    labelShowBg,
-    labelBgStyle,
-    labelBgPadding,
-    labelBgBorderRadius,
-    style,
-    source,
-    target,
-    isSelectable,
-    hidden,
-    sourceHandleId,
-    targetHandleId,
-    onContextMenu,
-    onMouseEnter,
-    onMouseMove,
-    onMouseLeave,
-    edgeUpdaterRadius,
-    onEdgeUpdate,
-    onEdgeUpdateStart,
-    onEdgeUpdateEnd,
-    markerEnd,
-    markerStart,
-    rfId,
-    ariaLabel,
-    isFocusable,
-    isUpdatable,
-    pathOptions,
-    interactionWidth,
-  }: WrapEdgeProps): JSX.Element | null => {
-    const edgeRef = useRef<SVGGElement>(null);
-    const [updateHover, setUpdateHover] = useState<boolean>(false);
-    const [updating, setUpdating] = useState<boolean>(false);
-    const store = useStoreApi();
-    const edgePosition = useStore(
-      useCallback(
-        (state) => {
-          const sourceNode = state.nodeLookup.get(source);
-          const targetNode = state.nodeLookup.get(target);
-
-          if (!sourceNode || !targetNode) {
-            return null;
-          }
-
-          return getEdgePosition({
-            id,
-            sourceNode,
-            targetNode,
-            sourceHandle: sourceHandleId || null,
-            targetHandle: targetHandleId || null,
-            connectionMode: state.connectionMode,
-            onError: state.onError,
-          });
-        },
-        [source, target]
-      ),
-      shallow
-    );
-
-    const markerStartUrl = useMemo(() => `url(#${getMarkerId(markerStart, rfId)})`, [markerStart, rfId]);
-    const markerEndUrl = useMemo(() => `url(#${getMarkerId(markerEnd, rfId)})`, [markerEnd, rfId]);
-
-    if (hidden || !edgePosition) {
-      return null;
-    }
-
-    const onEdgeClick = (event: React.MouseEvent<SVGGElement, MouseEvent>): void => {
-      const { edges, addSelectedEdges, unselectNodesAndEdges, multiSelectionActive } = store.getState();
-      const edge = edges.find((e) => e.id === id);
-
-      if (!edge) {
-        return;
-      }
-
-      if (isSelectable) {
-        store.setState({ nodesSelectionActive: false });
-
-        if (edge.selected && multiSelectionActive) {
-          unselectNodesAndEdges({ nodes: [], edges: [edge] });
-          edgeRef.current?.blur();
-        } else {
-          addSelectedEdges([id]);
-        }
-      }
-
-      if (onClick) {
-        onClick(event, edge);
-      }
-    };
-
-    const onEdgeDoubleClickHandler = getMouseHandler(id, store.getState, onEdgeDoubleClick);
-    const onEdgeContextMenu = getMouseHandler(id, store.getState, onContextMenu);
-    const onEdgeMouseEnter = getMouseHandler(id, store.getState, onMouseEnter);
-    const onEdgeMouseMove = getMouseHandler(id, store.getState, onMouseMove);
-    const onEdgeMouseLeave = getMouseHandler(id, store.getState, onMouseLeave);
-
-    const handleEdgeUpdater = (event: React.MouseEvent<SVGGElement, MouseEvent>, isSourceHandle: boolean) => {
-      // avoid triggering edge updater if mouse btn is not left
-      if (event.button !== 0) {
-        return;
-      }
-
-      const {
-        autoPanOnConnect,
-        domNode,
-        edges,
-        isValidConnection,
-        connectionMode,
-        connectionRadius,
-        lib,
-        onConnectStart,
-        onConnectEnd,
-        cancelConnection,
-        nodes,
-        panBy,
-        updateConnection,
-      } = store.getState();
-      const nodeId = isSourceHandle ? target : source;
-      const handleId = (isSourceHandle ? targetHandleId : sourceHandleId) || null;
-      const handleType = isSourceHandle ? 'target' : 'source';
-
-      const isTarget = isSourceHandle;
-      const edge = edges.find((e) => e.id === id)!;
-
-      setUpdating(true);
-      onEdgeUpdateStart?.(event, edge, handleType);
-
-      const _onEdgeUpdateEnd = (evt: MouseEvent | TouchEvent) => {
-        setUpdating(false);
-        onEdgeUpdateEnd?.(evt, edge, handleType);
-      };
-
-      const onConnectEdge = (connection: Connection) => onEdgeUpdate?.(edge, connection);
-
-      XYHandle.onPointerDown(event.nativeEvent, {
-        autoPanOnConnect,
-        connectionMode,
-        connectionRadius,
-        domNode,
-        handleId,
-        nodeId,
-        nodes,
-        isTarget,
-        edgeUpdaterType: handleType,
-        lib,
-        cancelConnection,
-        panBy,
-        isValidConnection,
-        onConnect: onConnectEdge,
-        onConnectStart,
-        onConnectEnd,
-        onEdgeUpdateEnd: _onEdgeUpdateEnd,
-        updateConnection,
-        getTransform: () => store.getState().transform,
-      });
-    };
-
-    const onEdgeUpdaterSourceMouseDown = (event: React.MouseEvent<SVGGElement, MouseEvent>): void =>
-      handleEdgeUpdater(event, true);
-    const onEdgeUpdaterTargetMouseDown = (event: React.MouseEvent<SVGGElement, MouseEvent>): void =>
-      handleEdgeUpdater(event, false);
-
-    const onEdgeUpdaterMouseEnter = () => setUpdateHover(true);
-    const onEdgeUpdaterMouseOut = () => setUpdateHover(false);
-
-    const inactive = !isSelectable && !onClick;
-
-    const onKeyDown = (event: KeyboardEvent) => {
-      if (elementSelectionKeys.includes(event.key) && isSelectable) {
-        const { unselectNodesAndEdges, addSelectedEdges, edges } = store.getState();
-        const unselect = event.key === 'Escape';
-
-        if (unselect) {
-          edgeRef.current?.blur();
-          unselectNodesAndEdges({ edges: [edges.find((e) => e.id === id)!] });
-        } else {
-          addSelectedEdges([id]);
-        }
-      }
-    };
-
-    return (
-      <g
-        className={cc([
-          'react-flow__edge',
-          `react-flow__edge-${type}`,
-          className,
-          { selected, animated, inactive, updating: updateHover },
-        ])}
-        onClick={onEdgeClick}
-        onDoubleClick={onEdgeDoubleClickHandler}
-        onContextMenu={onEdgeContextMenu}
-        onMouseEnter={onEdgeMouseEnter}
-        onMouseMove={onEdgeMouseMove}
-        onMouseLeave={onEdgeMouseLeave}
-        onKeyDown={isFocusable ? onKeyDown : undefined}
-        tabIndex={isFocusable ? 0 : undefined}
-        role={isFocusable ? 'button' : 'img'}
-        data-id={id}
-        data-testid={`rf__edge-${id}`}
-        aria-label={ariaLabel === null ? undefined : ariaLabel ? ariaLabel : `Edge from ${source} to ${target}`}
-        aria-describedby={isFocusable ? `${ARIA_EDGE_DESC_KEY}-${rfId}` : undefined}
-        ref={edgeRef}
-      >
-        {!updating && (
-          <EdgeComponent
-            id={id}
-            source={source}
-            target={target}
-            selected={selected}
-            animated={animated}
-            label={label}
-            labelStyle={labelStyle}
-            labelShowBg={labelShowBg}
-            labelBgStyle={labelBgStyle}
-            labelBgPadding={labelBgPadding}
-            labelBgBorderRadius={labelBgBorderRadius}
-            data={data}
-            style={style}
-            sourceX={edgePosition.sourceX}
-            sourceY={edgePosition.sourceY}
-            targetX={edgePosition.targetX}
-            targetY={edgePosition.targetY}
-            sourcePosition={edgePosition.sourcePosition}
-            targetPosition={edgePosition.targetPosition}
-            sourceHandleId={sourceHandleId}
-            targetHandleId={targetHandleId}
-            markerStart={markerStartUrl}
-            markerEnd={markerEndUrl}
-            pathOptions={pathOptions}
-            interactionWidth={interactionWidth}
-          />
-        )}
-        {isUpdatable && (
-          <>
-            {(isUpdatable === 'source' || isUpdatable === true) && (
-              <EdgeAnchor
-                position={edgePosition.sourcePosition}
-                centerX={edgePosition.sourceX}
-                centerY={edgePosition.sourceY}
-                radius={edgeUpdaterRadius}
-                onMouseDown={onEdgeUpdaterSourceMouseDown}
-                onMouseEnter={onEdgeUpdaterMouseEnter}
-                onMouseOut={onEdgeUpdaterMouseOut}
-                type="source"
-              />
-            )}
-            {(isUpdatable === 'target' || isUpdatable === true) && (
-              <EdgeAnchor
-                position={edgePosition.targetPosition}
-                centerX={edgePosition.targetX}
-                centerY={edgePosition.targetY}
-                radius={edgeUpdaterRadius}
-                onMouseDown={onEdgeUpdaterTargetMouseDown}
-                onMouseEnter={onEdgeUpdaterMouseEnter}
-                onMouseOut={onEdgeUpdaterMouseOut}
-                type="target"
-              />
-            )}
-          </>
-        )}
-      </g>
-    );
-  };
-
-  EdgeWrapper.displayName = 'EdgeWrapper';
-
-  return memo(EdgeWrapper);
-};

packages/react/src/components/NodeWrapper/index.tsx

@@ -1,4 +1,4 @@
-import { useEffect, useRef, memo, type MouseEvent, type KeyboardEvent, ComponentType } from 'react';
+import { useEffect, useRef, memo, type MouseEvent, type KeyboardEvent } from 'react';
 import cc from 'classcat';
 import {
   clampPosition,
@@ -7,7 +7,6 @@ import {
   getPositionWithOrigin,
   internalsSymbol,
   isInputDOMNode,
-  NodeProps,
 } from '@xyflow/system';
 
 import { useStore, useStoreApi } from '../../hooks/useStore';
@@ -17,7 +16,7 @@ import useDrag from '../../hooks/useDrag';
 import useUpdateNodePositions from '../../hooks/useUpdateNodePositions';
 import { handleNodeClick } from '../Nodes/utils';
 import type { NodeWrapperProps } from '../../types';
-import { arrowKeyDiffs } from './utils';
+import { arrowKeyDiffs, builtinNodeTypes } from './utils';
 
 const NodeWrapper = ({
   id,
@@ -44,10 +43,12 @@ const NodeWrapper = ({
   const node = useStore((s) => s.nodeLookup.get(id)!);
 
   let nodeType = node.type || 'default';
+  let NodeComponent = nodeTypes?.[nodeType] || builtinNodeTypes[nodeType];
 
-  if (!nodeTypes[nodeType]) {
+  if (NodeComponent === undefined) {
     onError?.('003', errorMessages['error003'](nodeType));
     nodeType = 'default';
+    NodeComponent = builtinNodeTypes.default;
   }
 
   const isDraggable = !!(node.draggable || (nodesDraggable && typeof node.draggable === 'undefined'));
@@ -105,7 +106,6 @@ const NodeWrapper = ({
     return null;
   }
 
-  const NodeComponent = (nodeTypes[nodeType] || nodeTypes.default) as ComponentType<NodeProps>;
   const width = node.width ?? undefined;
   const height = node.height ?? undefined;
   const computedWidth = node.computed?.width;

packages/react/src/components/NodeWrapper/utils.tsx

@@ -1,4 +1,11 @@
-import { XYPosition } from '@xyflow/system';
+import type { ComponentType } from 'react';
+import type { NodeProps, XYPosition } from '@xyflow/system';
+
+import InputNode from '../Nodes/InputNode';
+import DefaultNode from '../Nodes/DefaultNode';
+import GroupNode from '../Nodes/GroupNode';
+import OutputNode from '../Nodes/OutputNode';
+import type { NodeTypes } from '../../types';
 
 export const arrowKeyDiffs: Record<string, XYPosition> = {
   ArrowUp: { x: 0, y: -1 },
@@ -6,3 +13,10 @@ export const arrowKeyDiffs: Record<string, XYPosition> = {
   ArrowLeft: { x: -1, y: 0 },
   ArrowRight: { x: 1, y: 0 },
 };
+
+export const builtinNodeTypes: NodeTypes = {
+  input: InputNode as ComponentType<NodeProps>,
+  default: DefaultNode as ComponentType<NodeProps>,
+  output: OutputNode as ComponentType<NodeProps>,
+  group: GroupNode as ComponentType<NodeProps>,
+};

packages/react/src/container/EdgeRenderer/index.tsx

@@ -1,13 +1,13 @@
 import { memo, ReactNode } from 'react';
 import { shallow } from 'zustand/shallow';
 import cc from 'classcat';
-import { errorMessages } from '@xyflow/system';
 
 import { useStore } from '../../hooks/useStore';
 import useVisibleEdges from '../../hooks/useVisibleEdges';
 import MarkerDefinitions from './MarkerDefinitions';
 import { GraphViewProps } from '../GraphView';
-import type { EdgeTypesWrapped, ReactFlowState } from '../../types';
+import type { ReactFlowState } from '../../types';
+import EdgeWrapper from '../../components/EdgeWrapper';
 
 type EdgeRendererProps = Pick<
   GraphViewProps,
@@ -27,8 +27,8 @@ type EdgeRendererProps = Pick<
   | 'elevateEdgesOnSelect'
   | 'rfId'
   | 'disableKeyboardA11y'
+  | 'edgeTypes'
 > & {
-  edgeTypes: EdgeTypesWrapped;
   elevateEdgesOnSelect: boolean;
   children: ReactNode;
 };
@@ -74,14 +74,6 @@ const EdgeRenderer = ({
           {isMaxLevel && <MarkerDefinitions defaultColor={defaultMarkerColor} rfId={rfId} />}
           <>
             {edges.map((edge) => {
-              let edgeType = edge.type || 'default';
-
-              if (!edgeTypes[edgeType]) {
-                onError?.('011', errorMessages['error011'](edgeType));
-                edgeType = 'default';
-              }
-
-              const EdgeComponent = edgeTypes[edgeType];
               const isFocusable = !!(edge.focusable || (edgesFocusable && typeof edge.focusable === 'undefined'));
               const isUpdatable =
                 typeof onEdgeUpdate !== 'undefined' &&
@@ -92,7 +84,7 @@ const EdgeRenderer = ({
               );
 
               return (
-                <EdgeComponent
+                <EdgeWrapper
                   key={edge.id}
                   id={edge.id}
                   className={cc([edge.className, noPanClassName])}
@@ -131,6 +123,8 @@ const EdgeRenderer = ({
                   isUpdatable={isUpdatable}
                   pathOptions={'pathOptions' in edge ? edge.pathOptions : undefined}
                   interactionWidth={edge.interactionWidth}
+                  onError={onError}
+                  edgeTypes={edgeTypes}
                 />
               );
             })}

packages/react/src/container/EdgeRenderer/utils.ts

@@ -1,37 +0,0 @@
-import type { ComponentType } from 'react';
-
-import {
-  BezierEdgeInternal,
-  SmoothStepEdgeInternal,
-  StepEdgeInternal,
-  StraightEdgeInternal,
-  SimpleBezierEdgeInternal,
-} from '../../components/Edges';
-import wrapEdge from '../../components/Edges/wrapEdge';
-import type { EdgeProps, EdgeTypes, EdgeTypesWrapped } from '../../types';
-
-export type CreateEdgeTypes = (edgeTypes: EdgeTypes) => EdgeTypesWrapped;
-
-export function createEdgeTypes(edgeTypes: EdgeTypes): EdgeTypesWrapped {
-  const standardTypes: EdgeTypesWrapped = {
-    default: wrapEdge((edgeTypes.default || BezierEdgeInternal) as ComponentType<EdgeProps>),
-    straight: wrapEdge((edgeTypes.bezier || StraightEdgeInternal) as ComponentType<EdgeProps>),
-    step: wrapEdge((edgeTypes.step || StepEdgeInternal) as ComponentType<EdgeProps>),
-    smoothstep: wrapEdge((edgeTypes.step || SmoothStepEdgeInternal) as ComponentType<EdgeProps>),
-    simplebezier: wrapEdge((edgeTypes.simplebezier || SimpleBezierEdgeInternal) as ComponentType<EdgeProps>),
-  };
-
-  const wrappedTypes = {} as EdgeTypesWrapped;
-  const specialTypes: EdgeTypesWrapped = Object.keys(edgeTypes)
-    .filter((k) => !['default', 'bezier'].includes(k))
-    .reduce((res, key) => {
-      res[key] = wrapEdge((edgeTypes[key] || BezierEdgeInternal) as ComponentType<EdgeProps>);
-
-      return res;
-    }, wrappedTypes);
-
-  return {
-    ...standardTypes,
-    ...specialTypes,
-  };
-}

packages/react/src/container/GraphView/index.tsx

@@ -8,19 +8,15 @@ import useOnInitHandler from '../../hooks/useOnInitHandler';
 import useViewportSync from '../../hooks/useViewportSync';
 import ConnectionLine from '../../components/ConnectionLine';
 import type { ReactFlowProps } from '../../types';
-import { createNodeTypes } from '../NodeRenderer/utils';
+import useNodeOrEdgeTypesWarning from './useNodeOrEdgeTypesWarning';
-import { createEdgeTypes } from '../EdgeRenderer/utils';
-import { useNodeOrEdgeTypes } from './utils';
 
 export type GraphViewProps = Omit<
   ReactFlowProps,
-  'onSelectionChange' | 'nodes' | 'edges' | 'nodeTypes' | 'edgeTypes' | 'onMove' | 'onMoveStart' | 'onMoveEnd'
+  'onSelectionChange' | 'nodes' | 'edges' | 'onMove' | 'onMoveStart' | 'onMoveEnd'
 > &
   Required<
     Pick<
       ReactFlowProps,
-      | 'nodeTypes'
-      | 'edgeTypes'
       | 'selectionKeyCode'
       | 'deleteKeyCode'
       | 'multiSelectionKeyCode'
@@ -108,8 +104,8 @@ const GraphView = ({
   viewport,
   onViewportChange,
 }: GraphViewProps) => {
-  const nodeTypesWrapped = useNodeOrEdgeTypes(nodeTypes, createNodeTypes);
+  useNodeOrEdgeTypesWarning(nodeTypes);
-  const edgeTypesWrapped = useNodeOrEdgeTypes(edgeTypes, createEdgeTypes);
+  useNodeOrEdgeTypesWarning(edgeTypes);
 
   useOnInitHandler(onInit);
   useViewportSync(viewport);
@@ -154,7 +150,7 @@ const GraphView = ({
     >
       <ViewportWrapper>
         <EdgeRenderer
-          edgeTypes={edgeTypesWrapped}
+          edgeTypes={edgeTypes}
           onEdgeClick={onEdgeClick}
           onEdgeDoubleClick={onEdgeDoubleClick}
           onEdgeUpdate={onEdgeUpdate}
@@ -182,7 +178,7 @@ const GraphView = ({
         <div className="react-flow__edgelabel-renderer" />
 
         <NodeRenderer
-          nodeTypes={nodeTypesWrapped}
+          nodeTypes={nodeTypes}
           onNodeClick={onNodeClick}
           onNodeDoubleClick={onNodeDoubleClick}
           onNodeMouseEnter={onNodeMouseEnter}

packages/react/src/container/GraphView/useNodeOrEdgeTypesWarning.ts

@@ -0,0 +1,27 @@
+import { useEffect, useRef } from 'react';
+import { errorMessages } from '@xyflow/system';
+
+import type { EdgeTypes, NodeTypes } from '../../types';
+import { useStoreApi } from '../../hooks/useStore';
+
+const emptyTypes = {};
+
+/*
+ * This hook warns the user if node or edgeTypes change.
+ */
+export function useNodeOrEdgeTypesWarning(nodeOrEdgeTypes?: NodeTypes): void;
+export function useNodeOrEdgeTypesWarning(nodeOrEdgeTypes?: EdgeTypes): void;
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export default function useNodeOrEdgeTypesWarning(nodeOrEdgeTypes: any = emptyTypes): any {
+  const updateCount = useRef(0);
+  const store = useStoreApi();
+
+  useEffect(() => {
+    if (process.env.NODE_ENV === 'development') {
+      if (updateCount.current > 1) {
+        store.getState().onError?.('002', errorMessages['error002']());
+      }
+      updateCount.current += 1;
+    }
+  }, [nodeOrEdgeTypes]);
+}

packages/react/src/container/GraphView/utils.ts

@@ -1,31 +0,0 @@
-import { useMemo, useRef } from 'react';
-import { shallow } from 'zustand/shallow';
-import { errorMessages } from '@xyflow/system';
-
-import { CreateEdgeTypes } from '../EdgeRenderer/utils';
-import { CreateNodeTypes } from '../NodeRenderer/utils';
-import type { EdgeTypes, EdgeTypesWrapped, NodeTypes, NodeTypesWrapped } from '../../types';
-import { useStoreApi } from '../../hooks/useStore';
-
-export function useNodeOrEdgeTypes(nodeOrEdgeTypes: NodeTypes, createTypes: CreateNodeTypes): NodeTypesWrapped;
-export function useNodeOrEdgeTypes(nodeOrEdgeTypes: EdgeTypes, createTypes: CreateEdgeTypes): EdgeTypesWrapped;
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-export function useNodeOrEdgeTypes(nodeOrEdgeTypes: any, createTypes: any): any {
-  const typesKeysRef = useRef<string[] | null>(null);
-  const store = useStoreApi();
-
-  const typesParsed = useMemo(() => {
-    if (process.env.NODE_ENV === 'development') {
-      const typeKeys = Object.keys(nodeOrEdgeTypes);
-
-      if (shallow(typesKeysRef.current, typeKeys)) {
-        store.getState().onError?.('002', errorMessages['error002']());
-      }
-
-      typesKeysRef.current = typeKeys;
-    }
-    return createTypes(nodeOrEdgeTypes);
-  }, [nodeOrEdgeTypes]);
-
-  return typesParsed;
-}

packages/react/src/container/NodeRenderer/index.tsx

@@ -7,7 +7,7 @@ import { containerStyle } from '../../styles/utils';
 import { GraphViewProps } from '../GraphView';
 import type { ReactFlowState } from '../../types';
 import useResizeObserver from './useResizeObserver';
-import NodeWrapper from '../../components/NodeWrapper/NodeWrapper';
+import NodeWrapper from '../../components/NodeWrapper';
 
 export type NodeRendererProps = Pick<
   GraphViewProps,

packages/react/src/container/NodeRenderer/utils.ts

@@ -1,32 +0,0 @@
-import type { ComponentType } from 'react';
-import type { NodeProps } from '@xyflow/system';
-
-import DefaultNode from '../../components/Nodes/DefaultNode';
-import InputNode from '../../components/Nodes/InputNode';
-import OutputNode from '../../components/Nodes/OutputNode';
-import GroupNode from '../../components/Nodes/GroupNode';
-import type { NodeTypes } from '../../types';
-
-export type CreateNodeTypes = (nodeTypes: NodeTypes) => NodeTypes;
-
-export function createNodeTypes(nodeTypes: NodeTypes): NodeTypes {
-  const builtinTypes: NodeTypes = {
-    input: (nodeTypes.input || InputNode) as ComponentType<NodeProps>,
-    default: (nodeTypes.default || DefaultNode) as ComponentType<NodeProps>,
-    output: (nodeTypes.output || OutputNode) as ComponentType<NodeProps>,
-    group: (nodeTypes.group || GroupNode) as ComponentType<NodeProps>,
-  };
-
-  const userProvidedTypes = Object.keys(nodeTypes)
-    .filter((k) => !['input', 'default', 'output', 'group'].includes(k))
-    .reduce<NodeTypes>((res, key) => {
-      res[key] = (nodeTypes[key] || DefaultNode) as ComponentType<NodeProps>;
-
-      return res;
-    }, {});
-
-  return {
-    ...builtinTypes,
-    ...userProvidedTypes,
-  };
-}

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

@@ -12,40 +12,15 @@ import {
 } from '@xyflow/system';
 
 import Attribution from '../../components/Attribution';
-import {
+
-  BezierEdgeInternal,
-  SmoothStepEdgeInternal,
-  StepEdgeInternal,
-  StraightEdgeInternal,
-  SimpleBezierEdgeInternal,
-} from '../../components/Edges';
-import DefaultNode from '../../components/Nodes/DefaultNode';
-import InputNode from '../../components/Nodes/InputNode';
-import OutputNode from '../../components/Nodes/OutputNode';
-import GroupNode from '../../components/Nodes/GroupNode';
 import SelectionListener from '../../components/SelectionListener';
 import StoreUpdater from '../../components/StoreUpdater';
 import A11yDescriptions from '../../components/A11yDescriptions';
 import GraphView from '../GraphView';
 import Wrapper from './Wrapper';
-import type { EdgeTypes, NodeTypes, ReactFlowProps, ReactFlowRefType } from '../../types';
+import type { ReactFlowProps, ReactFlowRefType } from '../../types';
 import useColorModeClass from '../../hooks/useColorModeClass';
 
-const defaultNodeTypes: NodeTypes = {
-  input: InputNode,
-  default: DefaultNode,
-  output: OutputNode,
-  group: GroupNode,
-};
-
-const defaultEdgeTypes: EdgeTypes = {
-  default: BezierEdgeInternal,
-  straight: StraightEdgeInternal,
-  step: StepEdgeInternal,
-  smoothstep: SmoothStepEdgeInternal,
-  simplebezier: SimpleBezierEdgeInternal,
-};
-
 const initNodeOrigin: NodeOrigin = [0, 0];
 const initSnapGrid: [number, number] = [15, 15];
 const initDefaultViewport: Viewport = { x: 0, y: 0, zoom: 1 };
@@ -66,8 +41,8 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
       defaultNodes,
       defaultEdges,
       className,
-      nodeTypes = defaultNodeTypes,
+      nodeTypes,
-      edgeTypes = defaultEdgeTypes,
+      edgeTypes,
       onNodeClick,
       onEdgeClick,
       onInit,

packages/react/src/types/edges.ts

@@ -15,9 +15,10 @@ import type {
   EdgePosition,
   Optional,
   StepPathOptions,
+  OnError,
 } from '@xyflow/system';
 
-import { Node } from '.';
+import { EdgeTypes, Node } from '.';
 
 export type EdgeLabelOptions = {
   label?: string | ReactNode;
@@ -57,7 +58,7 @@ export type Edge<T = any> = DefaultEdge<T> | SmoothStepEdgeType<T> | BezierEdgeT
 
 export type EdgeMouseHandler = (event: ReactMouseEvent, edge: Edge) => void;
 
-export type WrapEdgeProps<T = any> = Omit<Edge<T>, 'sourceHandle' | 'targetHandle'> & {
+export type EdgeWrapperProps<T = any> = Omit<Edge<T>, 'sourceHandle' | 'targetHandle'> & {
   onClick?: EdgeMouseHandler;
   onEdgeDoubleClick?: EdgeMouseHandler;
   sourceHandleId?: string | null;
@@ -75,6 +76,8 @@ export type WrapEdgeProps<T = any> = Omit<Edge<T>, 'sourceHandle' | 'targetHandl
   isUpdatable: EdgeUpdatable;
   isSelectable: boolean;
   pathOptions?: BezierPathOptions | SmoothStepPathOptions;
+  edgeTypes?: EdgeTypes;
+  onError?: OnError;
 };
 
 export type DefaultEdgeOptions = DefaultEdgeOptionsBase<Edge>;
@@ -90,7 +93,7 @@ export type EdgeProps<T = any> = Pick<
   Edge<T>,
   'id' | 'animated' | 'data' | 'style' | 'selected' | 'source' | 'target'
 > &
-  Pick<WrapEdgeProps, 'sourceHandleId' | 'targetHandleId' | 'interactionWidth'> &
+  Pick<EdgeWrapperProps, 'sourceHandleId' | 'targetHandleId' | 'interactionWidth'> &
   EdgePosition &
   EdgeLabelOptions & {
     markerStart?: string;

packages/react/src/types/general.ts

@@ -1,5 +1,5 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
-import type { ComponentType, MemoExoticComponent } from 'react';
+import type { ComponentType } from 'react';
 import {
   FitViewParamsBase,
   FitViewOptionsBase,
@@ -14,7 +14,7 @@ import {
   XYPosition,
 } from '@xyflow/system';
 
-import type { NodeChange, EdgeChange, Node, WrapNodeProps, Edge, EdgeProps, WrapEdgeProps, ReactFlowInstance } from '.';
+import type { NodeChange, EdgeChange, Node, Edge, EdgeProps, ReactFlowInstance } from '.';
 
 export type OnNodesChange = (changes: NodeChange[]) => void;
 export type OnEdgesChange = (changes: EdgeChange[]) => void;
@@ -24,9 +24,7 @@ export type OnEdgesDelete = (edges: Edge[]) => void;
 export type OnDelete = (params: { nodes: Node[]; edges: Edge[] }) => void;
 
 export type NodeTypes = { [key: string]: ComponentType<NodeProps> };
-export type NodeTypesWrapped = { [key: string]: MemoExoticComponent<ComponentType<WrapNodeProps>> };
 export type EdgeTypes = { [key: string]: ComponentType<EdgeProps> };
-export type EdgeTypesWrapped = { [key: string]: MemoExoticComponent<ComponentType<WrapEdgeProps>> };
 
 export type UnselectNodesAndEdgesParams = {
   nodes?: Node[];

packages/react/src/types/nodes.ts

@@ -32,7 +32,7 @@ export type NodeWrapperProps = {
   noPanClassName: string;
   rfId: string;
   disableKeyboardA11y: boolean;
-  nodeTypes: NodeTypes;
+  nodeTypes?: NodeTypes;
   nodeExtent?: CoordinateExtent;
   nodeOrigin: NodeOrigin;
   onError?: OnError;