From 10f31e79e66419ce0280113caa50be5adb07de86 Mon Sep 17 00:00:00 2001
From: moklick <info@moritzklack.com>
Date: Thu, 7 Dec 2023 12:24:03 +0100
Subject: [PATCH] refactor(types): use custom node type as generic

---
 .../additional-components/MiniMap/types.ts    | 16 ++---
 packages/react/src/hooks/useNodes.ts          |  4 +-
 .../react/src/hooks/useNodesEdgesState.ts     | 12 ++--
 packages/react/src/hooks/useReactFlow.ts      | 59 ++++++++--------
 packages/react/src/types/changes.ts           | 16 ++---
 packages/react/src/types/edges.ts             | 34 +++++++++-
 packages/react/src/types/general.ts           | 13 ++--
 packages/react/src/types/instance.ts          | 67 ++++++++++---------
 packages/react/src/types/nodes.ts             | 21 +++---
 packages/react/src/utils/changes.ts           |  8 +--
 packages/svelte/src/lib/types/edges.ts        | 18 +++++
 packages/svelte/src/lib/types/nodes.ts        |  8 +--
 12 files changed, 160 insertions(+), 116 deletions(-)

diff --git a/packages/react/src/additional-components/MiniMap/types.ts b/packages/react/src/additional-components/MiniMap/types.ts
index 5667f7a8..3b24b08d 100644
--- a/packages/react/src/additional-components/MiniMap/types.ts
+++ b/packages/react/src/additional-components/MiniMap/types.ts
@@ -4,12 +4,12 @@ import type { PanelPosition, XYPosition } from '@xyflow/system';
 
 import type { Node } from '../../types';
 
-export type GetMiniMapNodeAttribute<NodeData = any> = (node: Node<NodeData>) => string;
+export type GetMiniMapNodeAttribute<NodeType extends Node = Node> = (node: NodeType) => string;
 
-export type MiniMapProps<NodeData = any> = Omit<HTMLAttributes<SVGSVGElement>, 'onClick'> & {
-  nodeColor?: string | GetMiniMapNodeAttribute<NodeData>;
-  nodeStrokeColor?: string | GetMiniMapNodeAttribute<NodeData>;
-  nodeClassName?: string | GetMiniMapNodeAttribute<NodeData>;
+export type MiniMapProps<NodeType extends Node = Node> = Omit<HTMLAttributes<SVGSVGElement>, 'onClick'> & {
+  nodeColor?: string | GetMiniMapNodeAttribute<NodeType>;
+  nodeStrokeColor?: string | GetMiniMapNodeAttribute<NodeType>;
+  nodeClassName?: string | GetMiniMapNodeAttribute<NodeType>;
   nodeBorderRadius?: number;
   nodeStrokeWidth?: number;
   nodeComponent?: ComponentType<MiniMapNodeProps>;
@@ -18,7 +18,7 @@ export type MiniMapProps<NodeData = any> = Omit<HTMLAttributes<SVGSVGElement>, '
   maskStrokeWidth?: number;
   position?: PanelPosition;
   onClick?: (event: MouseEvent, position: XYPosition) => void;
-  onNodeClick?: (event: MouseEvent, node: Node<NodeData>) => void;
+  onNodeClick?: (event: MouseEvent, node: NodeType) => void;
   pannable?: boolean;
   zoomable?: boolean;
   ariaLabel?: string | null;
@@ -27,8 +27,8 @@ export type MiniMapProps<NodeData = any> = Omit<HTMLAttributes<SVGSVGElement>, '
   offsetScale?: number;
 };
 
-export type MiniMapNodes = Pick<
-  MiniMapProps,
+export type MiniMapNodes<NodeType extends Node = Node> = Pick<
+  MiniMapProps<NodeType>,
   'nodeColor' | 'nodeStrokeColor' | 'nodeClassName' | 'nodeBorderRadius' | 'nodeStrokeWidth' | 'nodeComponent'
 > & {
   onClick?: (event: MouseEvent, nodeId: string) => void;
diff --git a/packages/react/src/hooks/useNodes.ts b/packages/react/src/hooks/useNodes.ts
index cc75f142..1ccd1d85 100644
--- a/packages/react/src/hooks/useNodes.ts
+++ b/packages/react/src/hooks/useNodes.ts
@@ -5,8 +5,8 @@ import type { Node, ReactFlowState } from '../types';
 
 const nodesSelector = (state: ReactFlowState) => state.nodes;
 
-function useNodes<NodeData>(): Node<NodeData>[] {
-  const nodes = useStore(nodesSelector, shallow);
+function useNodes<NodeType extends Node = Node>(): NodeType[] {
+  const nodes = useStore(nodesSelector, shallow) as NodeType[];
 
   return nodes;
 }
diff --git a/packages/react/src/hooks/useNodesEdgesState.ts b/packages/react/src/hooks/useNodesEdgesState.ts
index 19754aa4..5255c7dd 100644
--- a/packages/react/src/hooks/useNodesEdgesState.ts
+++ b/packages/react/src/hooks/useNodesEdgesState.ts
@@ -11,14 +11,14 @@ type OnChange<ChangesType> = (changes: ChangesType[]) => void;
 // const [nodes, setNodes, onNodesChange] = useNodesState(initialNodes);
 function createUseItemsState(
   applyChanges: ApplyChanges<Node, NodeChange>
-): <NodeData = any>(
-  initialItems: Node<NodeData>[]
-) => [Node<NodeData>[], Dispatch<SetStateAction<Node<NodeData>[]>>, OnChange<NodeChange>];
+): <NodeType extends Node = Node>(
+  initialItems: NodeType[]
+) => [NodeType[], Dispatch<SetStateAction<NodeType[]>>, OnChange<NodeChange>];
 function createUseItemsState(
   applyChanges: ApplyChanges<Edge, EdgeChange>
-): <EdgeData = any>(
-  initialItems: Edge<EdgeData>[]
-) => [Edge<EdgeData>[], Dispatch<SetStateAction<Edge<EdgeData>[]>>, OnChange<EdgeChange>];
+): <EdgeType extends Edge = Edge>(
+  initialItems: EdgeType[]
+) => [EdgeType[], Dispatch<SetStateAction<EdgeType[]>>, OnChange<EdgeChange>];
 function createUseItemsState(
   applyChanges: ApplyChanges<any, any>
 ): (initialItems: any[]) => [any[], Dispatch<SetStateAction<any[]>>, OnChange<any>] {
diff --git a/packages/react/src/hooks/useReactFlow.ts b/packages/react/src/hooks/useReactFlow.ts
index 44e29c14..2875ca03 100644
--- a/packages/react/src/hooks/useReactFlow.ts
+++ b/packages/react/src/hooks/useReactFlow.ts
@@ -26,31 +26,34 @@ import type {
 } from '../types';
 
 /* eslint-disable-next-line @typescript-eslint/no-explicit-any */
-export default function useReactFlow<NodeData = any, EdgeData = any>(): ReactFlowInstance<NodeData, EdgeData> {
+export default function useReactFlow<NodeType extends Node = Node, EdgeType extends Edge = Edge>(): ReactFlowInstance<
+  NodeType,
+  EdgeType
+> {
   const viewportHelper = useViewportHelper();
   const store = useStoreApi();
 
-  const getNodes = useCallback<Instance.GetNodes<NodeData>>(() => {
-    return store.getState().nodes.map((n) => ({ ...n }));
+  const getNodes = useCallback<Instance.GetNodes<NodeType>>(() => {
+    return store.getState().nodes.map((n) => ({ ...n })) as NodeType[];
   }, []);
 
-  const getNode = useCallback<Instance.GetNode<NodeData>>((id) => {
-    return store.getState().nodeLookup.get(id);
+  const getNode = useCallback<Instance.GetNode<NodeType>>((id) => {
+    return store.getState().nodeLookup.get(id) as NodeType;
   }, []);
 
-  const getEdges = useCallback<Instance.GetEdges<EdgeData>>(() => {
+  const getEdges = useCallback<Instance.GetEdges<EdgeType>>(() => {
     const { edges = [] } = store.getState();
-    return edges.map((e) => ({ ...e }));
+    return edges.map((e) => ({ ...e })) as EdgeType[];
   }, []);
 
-  const getEdge = useCallback<Instance.GetEdge<EdgeData>>((id) => {
+  const getEdge = useCallback<Instance.GetEdge<EdgeType>>((id) => {
     const { edges = [] } = store.getState();
-    return edges.find((e) => e.id === id);
+    return edges.find((e) => e.id === id) as EdgeType;
   }, []);
 
-  const setNodes = useCallback<Instance.SetNodes<NodeData>>((payload) => {
+  const setNodes = useCallback<Instance.SetNodes<NodeType>>((payload) => {
     const { nodes, setNodes, hasDefaultNodes, onNodesChange } = store.getState();
-    const nextNodes = typeof payload === 'function' ? payload(nodes) : payload;
+    const nextNodes = typeof payload === 'function' ? payload(nodes as NodeType[]) : payload;
 
     if (hasDefaultNodes) {
       setNodes(nextNodes);
@@ -58,14 +61,14 @@ export default function useReactFlow<NodeData = any, EdgeData = any>(): ReactFlo
       const changes =
         nextNodes.length === 0
           ? nodes.map((node) => ({ type: 'remove', id: node.id } as NodeRemoveChange))
-          : nextNodes.map((node) => ({ item: node, type: 'reset' } as NodeResetChange<NodeData>));
+          : nextNodes.map((node) => ({ item: node, type: 'reset' } as NodeResetChange<NodeType>));
       onNodesChange(changes);
     }
   }, []);
 
-  const setEdges = useCallback<Instance.SetEdges<EdgeData>>((payload) => {
+  const setEdges = useCallback<Instance.SetEdges<EdgeType>>((payload) => {
     const { edges = [], setEdges, hasDefaultEdges, onEdgesChange } = store.getState();
-    const nextEdges = typeof payload === 'function' ? payload(edges) : payload;
+    const nextEdges = typeof payload === 'function' ? payload(edges as EdgeType[]) : payload;
 
     if (hasDefaultEdges) {
       setEdges(nextEdges);
@@ -73,12 +76,12 @@ export default function useReactFlow<NodeData = any, EdgeData = any>(): ReactFlo
       const changes =
         nextEdges.length === 0
           ? edges.map((edge) => ({ type: 'remove', id: edge.id } as EdgeRemoveChange))
-          : nextEdges.map((edge) => ({ item: edge, type: 'reset' } as EdgeResetChange<EdgeData>));
+          : nextEdges.map((edge) => ({ item: edge, type: 'reset' } as EdgeResetChange<EdgeType>));
       onEdgesChange(changes);
     }
   }, []);
 
-  const addNodes = useCallback<Instance.AddNodes<NodeData>>((payload) => {
+  const addNodes = useCallback<Instance.AddNodes<NodeType>>((payload) => {
     const nodes = Array.isArray(payload) ? payload : [payload];
     const { nodes: currentNodes, hasDefaultNodes, onNodesChange, setNodes } = store.getState();
 
@@ -86,29 +89,29 @@ export default function useReactFlow<NodeData = any, EdgeData = any>(): ReactFlo
       const nextNodes = [...currentNodes, ...nodes];
       setNodes(nextNodes);
     } else if (onNodesChange) {
-      const changes = nodes.map((node) => ({ item: node, type: 'add' } as NodeAddChange<NodeData>));
+      const changes = nodes.map((node) => ({ item: node, type: 'add' } as NodeAddChange<NodeType>));
       onNodesChange(changes);
     }
   }, []);
 
-  const addEdges = useCallback<Instance.AddEdges<EdgeData>>((payload) => {
+  const addEdges = useCallback<Instance.AddEdges<EdgeType>>((payload) => {
     const nextEdges = Array.isArray(payload) ? payload : [payload];
     const { edges = [], setEdges, hasDefaultEdges, onEdgesChange } = store.getState();
 
     if (hasDefaultEdges) {
       setEdges([...edges, ...nextEdges]);
     } else if (onEdgesChange) {
-      const changes = nextEdges.map((edge) => ({ item: edge, type: 'add' } as EdgeAddChange<EdgeData>));
+      const changes = nextEdges.map((edge) => ({ item: edge, type: 'add' } as EdgeAddChange<EdgeType>));
       onEdgesChange(changes);
     }
   }, []);
 
-  const toObject = useCallback<Instance.ToObject<NodeData, EdgeData>>(() => {
+  const toObject = useCallback<Instance.ToObject<NodeType, EdgeType>>(() => {
     const { nodes = [], edges = [], transform } = store.getState();
     const [x, y, zoom] = transform;
     return {
-      nodes: nodes.map((n) => ({ ...n })),
-      edges: edges.map((e) => ({ ...e })),
+      nodes: nodes.map((n) => ({ ...n })) as NodeType[],
+      edges: edges.map((e) => ({ ...e })) as EdgeType[],
       viewport: {
         x,
         y,
@@ -180,11 +183,9 @@ export default function useReactFlow<NodeData = any, EdgeData = any>(): ReactFlo
   }, []);
 
   const getNodeRect = useCallback(
-    (
-      nodeOrRect: Node<NodeData> | { id: Node['id'] } | Rect
-    ): [Rect | null, Node<NodeData> | null | undefined, boolean] => {
+    (nodeOrRect: NodeType | { id: Node['id'] } | Rect): [Rect | null, NodeType | null | undefined, boolean] => {
       const isRect = isRectObject(nodeOrRect);
-      const node = isRect ? null : store.getState().nodeLookup.get(nodeOrRect.id);
+      const node = isRect ? null : (store.getState().nodeLookup.get(nodeOrRect.id) as NodeType);
 
       if (!isRect && !node) {
         [null, null, isRect];
@@ -197,7 +198,7 @@ export default function useReactFlow<NodeData = any, EdgeData = any>(): ReactFlo
     []
   );
 
-  const getIntersectingNodes = useCallback<Instance.GetIntersectingNodes<NodeData>>(
+  const getIntersectingNodes = useCallback<Instance.GetIntersectingNodes<NodeType>>(
     (nodeOrRect, partially = true, nodes) => {
       const [nodeRect, node, isRect] = getNodeRect(nodeOrRect);
 
@@ -215,12 +216,12 @@ export default function useReactFlow<NodeData = any, EdgeData = any>(): ReactFlo
         const partiallyVisible = partially && overlappingArea > 0;
 
         return partiallyVisible || overlappingArea >= nodeRect.width * nodeRect.height;
-      });
+      }) as NodeType[];
     },
     []
   );
 
-  const isNodeIntersecting = useCallback<Instance.IsNodeIntersecting<NodeData>>(
+  const isNodeIntersecting = useCallback<Instance.IsNodeIntersecting<NodeType>>(
     (nodeOrRect, area, partially = true) => {
       const [nodeRect] = getNodeRect(nodeOrRect);
 
diff --git a/packages/react/src/types/changes.ts b/packages/react/src/types/changes.ts
index 69ee0f1f..440a4c6b 100644
--- a/packages/react/src/types/changes.ts
+++ b/packages/react/src/types/changes.ts
@@ -30,13 +30,13 @@ export type NodeRemoveChange = {
   type: 'remove';
 };
 
-export type NodeAddChange<NodeData = any> = {
-  item: Node<NodeData>;
+export type NodeAddChange<NodeType extends Node = Node> = {
+  item: NodeType;
   type: 'add';
 };
 
-export type NodeResetChange<NodeData = any> = {
-  item: Node<NodeData>;
+export type NodeResetChange<NodeType extends Node = Node> = {
+  item: NodeType;
   type: 'reset';
 };
 
@@ -50,12 +50,12 @@ export type NodeChange =
 
 export type EdgeSelectionChange = NodeSelectionChange;
 export type EdgeRemoveChange = NodeRemoveChange;
-export type EdgeAddChange<EdgeData = any> = {
-  item: Edge<EdgeData>;
+export type EdgeAddChange<EdgeType extends Edge = Edge> = {
+  item: EdgeType;
   type: 'add';
 };
-export type EdgeResetChange<EdgeData = any> = {
-  item: Edge<EdgeData>;
+export type EdgeResetChange<EdgeType extends Edge = Edge> = {
+  item: EdgeType;
   type: 'reset';
 };
 export type EdgeChange = EdgeSelectionChange | EdgeRemoveChange | EdgeAddChange | EdgeResetChange;
diff --git a/packages/react/src/types/edges.ts b/packages/react/src/types/edges.ts
index fecd7c7b..836087fc 100644
--- a/packages/react/src/types/edges.ts
+++ b/packages/react/src/types/edges.ts
@@ -1,6 +1,13 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
-import type { CSSProperties, HTMLAttributes, ReactNode, MouseEvent as ReactMouseEvent, ComponentType } from 'react';
+import type {
+  CSSProperties,
+  HTMLAttributes,
+  ReactNode,
+  MouseEvent as ReactMouseEvent,
+  ComponentType,
+  MemoExoticComponent,
+} from 'react';
 import type {
   EdgeBase,
   BezierPathOptions,
@@ -119,6 +126,9 @@ export type BaseEdgeProps = EdgeLabelOptions & {
   style?: CSSProperties;
 };
 
+/**
+ * Helper type for edge components that get exported by the library.
+ */
 export type EdgeComponentProps = EdgePosition &
   EdgeLabelOptions & {
     id?: EdgeProps['id'];
@@ -134,10 +144,29 @@ export type EdgeComponentWithPathOptions<PathOptions> = EdgeComponentProps & {
   pathOptions?: PathOptions;
 };
 
+/**
+ * BezierEdge component props
+ */
 export type BezierEdgeProps = EdgeComponentWithPathOptions<BezierPathOptions>;
+
+/**
+ * SmoothStepEdge component props
+ */
 export type SmoothStepEdgeProps = EdgeComponentWithPathOptions<SmoothStepPathOptions>;
+
+/**
+ * StepEdge component props
+ */
 export type StepEdgeProps = EdgeComponentWithPathOptions<StepPathOptions>;
+
+/**
+ * StraightEdge component props
+ */
 export type StraightEdgeProps = Omit<EdgeComponentProps, 'sourcePosition' | 'targetPosition'>;
+
+/**
+ * SimpleBezier component props
+ */
 export type SimpleBezierEdgeProps = EdgeComponentProps;
 
 export type OnEdgeUpdateFunc<T = any> = (oldEdge: Edge<T>, newConnection: Connection) => void;
@@ -157,3 +186,6 @@ export type ConnectionLineComponentProps = {
 };
 
 export type ConnectionLineComponent = ComponentType<ConnectionLineComponentProps>;
+
+export type EdgeTypes = { [key: string]: ComponentType<EdgeProps> };
+export type EdgeTypesWrapped = { [key: string]: MemoExoticComponent<ComponentType<WrapEdgeProps>> };
diff --git a/packages/react/src/types/general.ts b/packages/react/src/types/general.ts
index b2e7098a..a212b057 100644
--- a/packages/react/src/types/general.ts
+++ b/packages/react/src/types/general.ts
@@ -1,9 +1,7 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
-import type { ComponentType, MemoExoticComponent } from 'react';
 import {
   FitViewParamsBase,
   FitViewOptionsBase,
-  NodeProps,
   ZoomInOut,
   ZoomTo,
   SetViewport,
@@ -14,7 +12,7 @@ import {
   XYPosition,
 } from '@xyflow/system';
 
-import type { NodeChange, EdgeChange, Node, WrapNodeProps, Edge, EdgeProps, WrapEdgeProps, ReactFlowInstance } from '.';
+import type { NodeChange, EdgeChange, Node, Edge, ReactFlowInstance } from '.';
 
 export type OnNodesChange = (changes: NodeChange[]) => void;
 export type OnEdgesChange = (changes: EdgeChange[]) => void;
@@ -23,11 +21,6 @@ export type OnNodesDelete = (nodes: Node[]) => void;
 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[];
   edges?: Edge[];
@@ -43,7 +36,9 @@ export type OnSelectionChangeFunc = (params: OnSelectionChangeParams) => void;
 export type FitViewParams = FitViewParamsBase<Node>;
 export type FitViewOptions = FitViewOptionsBase<Node>;
 export type FitView = (fitViewOptions?: FitViewOptions) => boolean;
-export type OnInit<NodeData = any, EdgeData = any> = (reactFlowInstance: ReactFlowInstance<NodeData, EdgeData>) => void;
+export type OnInit<NodeType extends Node = Node, EdgeType extends Edge = Edge> = (
+  reactFlowInstance: ReactFlowInstance<NodeType, EdgeType>
+) => void;
 
 export type ViewportHelperFunctions = {
   zoomIn: ZoomInOut;
diff --git a/packages/react/src/types/instance.ts b/packages/react/src/types/instance.ts
index f229ec87..6ac81e04 100644
--- a/packages/react/src/types/instance.ts
+++ b/packages/react/src/types/instance.ts
@@ -3,9 +3,9 @@
 import type { Rect, Viewport } from '@xyflow/system';
 import type { Node, Edge, ViewportHelperFunctions } from '.';
 
-export type ReactFlowJsonObject<NodeData = any, EdgeData = any> = {
-  nodes: Node<NodeData>[];
-  edges: Edge<EdgeData>[];
+export type ReactFlowJsonObject<NodeType extends Node = Node, EdgeType extends Edge = Edge> = {
+  nodes: NodeType[];
+  edges: EdgeType[];
   viewport: Viewport;
 };
 
@@ -15,30 +15,33 @@ export type DeleteElementsOptions = {
 };
 
 export namespace Instance {
-  export type GetNodes<NodeData> = () => Node<NodeData>[];
-  export type SetNodes<NodeData> = (
-    payload: Node<NodeData>[] | ((nodes: Node<NodeData>[]) => Node<NodeData>[])
+  export type GetNodes<NodeType extends Node = Node> = () => NodeType[];
+  export type SetNodes<NodeType extends Node = Node> = (
+    payload: NodeType[] | ((nodes: NodeType[]) => NodeType[])
   ) => void;
-  export type AddNodes<NodeData> = (payload: Node<NodeData>[] | Node<NodeData>) => void;
-  export type GetNode<NodeData> = (id: string) => Node<NodeData> | undefined;
-  export type GetEdges<EdgeData> = () => Edge<EdgeData>[];
-  export type SetEdges<EdgeData> = (
-    payload: Edge<EdgeData>[] | ((edges: Edge<EdgeData>[]) => Edge<EdgeData>[])
+  export type AddNodes<NodeType extends Node = Node> = (payload: NodeType[] | NodeType) => void;
+  export type GetNode<NodeType extends Node = Node> = (id: string) => NodeType | undefined;
+  export type GetEdges<EdgeType extends Edge = Edge> = () => EdgeType[];
+  export type SetEdges<EdgeType extends Edge = Edge> = (
+    payload: EdgeType[] | ((edges: EdgeType[]) => EdgeType[])
   ) => void;
-  export type GetEdge<EdgeData> = (id: string) => Edge<EdgeData> | undefined;
-  export type AddEdges<EdgeData> = (payload: Edge<EdgeData>[] | Edge<EdgeData>) => void;
-  export type ToObject<NodeData = any, EdgeData = any> = () => ReactFlowJsonObject<NodeData, EdgeData>;
+  export type GetEdge<EdgeType extends Edge = Edge> = (id: string) => EdgeType | undefined;
+  export type AddEdges<EdgeType extends Edge = Edge> = (payload: EdgeType[] | EdgeType) => void;
+  export type ToObject<NodeType extends Node = Node, EdgeType extends Edge = Edge> = () => ReactFlowJsonObject<
+    NodeType,
+    EdgeType
+  >;
   export type DeleteElements = ({ nodes, edges }: DeleteElementsOptions) => {
     deletedNodes: Node[];
     deletedEdges: Edge[];
   };
-  export type GetIntersectingNodes<NodeData> = (
-    node: Node<NodeData> | { id: Node['id'] } | Rect,
+  export type GetIntersectingNodes<NodeType extends Node = Node> = (
+    node: NodeType | { id: Node['id'] } | Rect,
     partially?: boolean,
-    nodes?: Node<NodeData>[]
-  ) => Node<NodeData>[];
-  export type IsNodeIntersecting<NodeData> = (
-    node: Node<NodeData> | { id: Node['id'] } | Rect,
+    nodes?: NodeType[]
+  ) => NodeType[];
+  export type IsNodeIntersecting<NodeType extends Node = Node> = (
+    node: NodeType | { id: Node['id'] } | Rect,
     area: Rect,
     partially?: boolean
   ) => boolean;
@@ -47,18 +50,18 @@ export namespace Instance {
   export type getOutgoers = (node: string | Node | { id: Node['id'] }) => Node[];
 }
 
-export type ReactFlowInstance<NodeData = any, EdgeData = any> = {
-  getNodes: Instance.GetNodes<NodeData>;
-  setNodes: Instance.SetNodes<NodeData>;
-  addNodes: Instance.AddNodes<NodeData>;
-  getNode: Instance.GetNode<NodeData>;
-  getEdges: Instance.GetEdges<EdgeData>;
-  setEdges: Instance.SetEdges<EdgeData>;
-  addEdges: Instance.AddEdges<EdgeData>;
-  getEdge: Instance.GetEdge<EdgeData>;
-  toObject: Instance.ToObject<NodeData, EdgeData>;
+export type ReactFlowInstance<NodeType extends Node = Node, EdgeType extends Edge = Edge> = {
+  getNodes: Instance.GetNodes<NodeType>;
+  setNodes: Instance.SetNodes<NodeType>;
+  addNodes: Instance.AddNodes<NodeType>;
+  getNode: Instance.GetNode<NodeType>;
+  getEdges: Instance.GetEdges<EdgeType>;
+  setEdges: Instance.SetEdges<EdgeType>;
+  addEdges: Instance.AddEdges<EdgeType>;
+  getEdge: Instance.GetEdge<EdgeType>;
+  toObject: Instance.ToObject<NodeType, EdgeType>;
   deleteElements: Instance.DeleteElements;
-  getIntersectingNodes: Instance.GetIntersectingNodes<NodeData>;
-  isNodeIntersecting: Instance.IsNodeIntersecting<NodeData>;
+  getIntersectingNodes: Instance.GetIntersectingNodes<NodeType>;
+  isNodeIntersecting: Instance.IsNodeIntersecting<NodeType>;
   viewportInitialized: boolean;
 } & Omit<ViewportHelperFunctions, 'initialized'>;
diff --git a/packages/react/src/types/nodes.ts b/packages/react/src/types/nodes.ts
index 4d10b76b..7bccfe0a 100644
--- a/packages/react/src/types/nodes.ts
+++ b/packages/react/src/types/nodes.ts
@@ -1,5 +1,6 @@
-import type { CSSProperties, MouseEvent as ReactMouseEvent } from 'react';
-import type { NodeBase, XYPosition } from '@xyflow/system';
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import type { CSSProperties, ComponentType, MemoExoticComponent, MouseEvent as ReactMouseEvent } from 'react';
+import type { NodeBase, NodeProps, XYPosition } from '@xyflow/system';
 
 /**
  * The node data structure that gets used for the nodes prop.
@@ -9,17 +10,8 @@ export type Node<NodeData = any, NodeType extends string | undefined = string |
   NodeData,
   NodeType
 > & {
-  /**
-   * Inline style object
-   */
   style?: CSSProperties;
-  /**
-   * Inline style object
-   */
   className?: string;
-  /**
-   * Inline style object
-   */
   resizing?: boolean;
 };
 
@@ -27,8 +19,8 @@ export type NodeMouseHandler = (event: ReactMouseEvent, node: Node) => void;
 export type NodeDragHandler = (event: ReactMouseEvent, node: Node, nodes: Node[]) => void;
 export type SelectionDragHandler = (event: ReactMouseEvent, nodes: Node[]) => void;
 
-export type WrapNodeProps<NodeData = any> = Pick<
-  Node<NodeData>,
+export type WrapNodeProps<NodeType extends Node = Node> = Pick<
+  NodeType,
   'id' | 'data' | 'style' | 'className' | 'dragHandle' | 'sourcePosition' | 'targetPosition' | 'hidden' | 'ariaLabel'
 > & {
   type: string;
@@ -59,3 +51,6 @@ export type WrapNodeProps<NodeData = any> = Pick<
   width?: number;
   height?: number;
 };
+
+export type NodeTypes = { [key: string]: ComponentType<NodeProps> };
+export type NodeTypesWrapped = { [key: string]: MemoExoticComponent<ComponentType<WrapNodeProps>> };
diff --git a/packages/react/src/utils/changes.ts b/packages/react/src/utils/changes.ts
index 0cf87e6f..72e28b9e 100644
--- a/packages/react/src/utils/changes.ts
+++ b/packages/react/src/utils/changes.ts
@@ -159,8 +159,8 @@ function applyChanges(changes: any[], elements: any[]): any[] {
       <ReactFLow nodes={nodes} edges={edges} onNodesChange={onNodesChange} />
     );
  */
-export function applyNodeChanges<NodeData = any>(changes: NodeChange[], nodes: Node<NodeData>[]): Node<NodeData>[] {
-  return applyChanges(changes, nodes) as Node<NodeData>[];
+export function applyNodeChanges<NodeType extends Node = Node>(changes: NodeChange[], nodes: NodeType[]): NodeType[] {
+  return applyChanges(changes, nodes) as NodeType[];
 }
 
 /**
@@ -183,8 +183,8 @@ export function applyNodeChanges<NodeData = any>(changes: NodeChange[], nodes: N
       <ReactFLow nodes={nodes} edges={edges} onEdgesChange={onEdgesChange} />
     );
  */
-export function applyEdgeChanges<EdgeData = any>(changes: EdgeChange[], edges: Edge<EdgeData>[]): Edge<EdgeData>[] {
-  return applyChanges(changes, edges) as Edge<EdgeData>[];
+export function applyEdgeChanges<EdgeType extends Edge = Edge>(changes: EdgeChange[], edges: EdgeType[]): EdgeType[] {
+  return applyChanges(changes, edges) as EdgeType[];
 }
 
 export const createSelectionChange = (id: string, selected: boolean) => ({
diff --git a/packages/svelte/src/lib/types/edges.ts b/packages/svelte/src/lib/types/edges.ts
index c1d3778d..218c7e6f 100644
--- a/packages/svelte/src/lib/types/edges.ts
+++ b/packages/svelte/src/lib/types/edges.ts
@@ -53,6 +53,9 @@ export type EdgeProps<T = any> = Omit<Edge<T>, 'sourceHandle' | 'targetHandle' |
     targetHandleId?: string | null;
   };
 
+/**
+ * Helper type for edge components that get exported by the library.
+ */
 export type EdgeComponentProps = EdgePosition & {
   id?: EdgeProps['id'];
   hidden?: EdgeProps['hidden'];
@@ -73,9 +76,24 @@ export type EdgeComponentWithPathOptions<PathOptions> = EdgeComponentProps & {
   pathOptions?: PathOptions;
 };
 
+/**
+ * BezierEdge component props
+ */
 export type BezierEdgeProps = EdgeComponentWithPathOptions<BezierPathOptions>;
+
+/**
+ * SmoothStepEdge component props
+ */
 export type SmoothStepEdgeProps = EdgeComponentWithPathOptions<SmoothStepPathOptions>;
+
+/**
+ * StepEdge component props
+ */
 export type StepEdgeProps = EdgeComponentWithPathOptions<StepPathOptions>;
+
+/**
+ * StraightEdge component props
+ */
 export type StraightEdgeProps = Omit<EdgeComponentProps, 'sourcePosition' | 'targetPosition'>;
 
 export type EdgeTypes = Record<string, ComponentType<SvelteComponent<EdgeProps>>>;
diff --git a/packages/svelte/src/lib/types/nodes.ts b/packages/svelte/src/lib/types/nodes.ts
index 292ae4fa..4121754a 100644
--- a/packages/svelte/src/lib/types/nodes.ts
+++ b/packages/svelte/src/lib/types/nodes.ts
@@ -1,11 +1,11 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
-
 import type { ComponentType, SvelteComponent } from 'svelte';
 import type { NodeBase, NodeProps } from '@xyflow/system';
 
-// @todo: currently the helper function only like Node from '@reactflow/core'
-// we need a base node type or helpes that accept Node like types
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
+/**
+ * The node data structure that gets used for the nodes prop.
+ * @public
+ */
 export type Node<
   NodeData = any,
   NodeType extends string | undefined = string | undefined