refactor(types): use custom node type as generic

2023-12-07 12:24:03 +01:00
parent 0f657a5193
commit 10f31e79e6
12 changed files with 160 additions and 116 deletions
--- 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;
--- 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;
 }
--- 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>] {
--- 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);

--- 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;
--- 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>> };
--- 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;
--- 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'>;
--- 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>> };
--- 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) => ({
--- 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>>>;
--- 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