refactor(react): add generics to useStoreApi

packages/react/src/additional-components/MiniMap/MiniMap.tsx

@@ -63,7 +63,7 @@ function MiniMapComponent<NodeType extends Node = Node>({
   zoomStep = 10,
   offsetScale = 5,
 }: MiniMapProps<NodeType>) {
-  const store = useStoreApi();
+  const store = useStoreApi<NodeType>();
   const svg = useRef<SVGSVGElement>(null);
   const { boundingRect, viewBB, rfId, panZoom, translateExtent, flowWidth, flowHeight } = useStore(selector, shallow);
   const elementWidth = (style?.width as number) ?? defaultWidth;
@@ -120,7 +120,7 @@ function MiniMapComponent<NodeType extends Node = Node>({
 
   const onSvgNodeClick = onNodeClick
     ? useCallback((event: MouseEvent, nodeId: string) => {
-        const node = store.getState().nodeLookup.get(nodeId)! as NodeType;
+        const node = store.getState().nodeLookup.get(nodeId)!;
         onNodeClick(event, node);
       }, [])
     : undefined;

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

@@ -2,7 +2,6 @@
  * The nodes selection rectangle gets displayed when a user
  * made a selection with on or several nodes
  */
-
 import { useRef, useEffect, type MouseEvent, type KeyboardEvent } from 'react';
 import cc from 'classcat';
 import { shallow } from 'zustand/shallow';
@@ -37,7 +36,7 @@ export function NodesSelection<NodeType extends Node>({
   noPanClassName,
   disableKeyboardA11y,
 }: NodesSelectionProps<NodeType>) {
-  const store = useStoreApi();
+  const store = useStoreApi<NodeType>();
   const { width, height, transformString, userSelectionActive } = useStore(selector, shallow);
   const moveSelectedNodes = useMoveSelectedNodes();
 
@@ -62,7 +61,7 @@ export function NodesSelection<NodeType extends Node>({
   const onContextMenu = onSelectionContextMenu
     ? (event: MouseEvent) => {
         const selectedNodes = store.getState().nodes.filter((n) => n.selected);
-        onSelectionContextMenu(event, selectedNodes as NodeType[]);
+        onSelectionContextMenu(event, selectedNodes);
       }
     : undefined;
 

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

@@ -115,7 +115,7 @@ export function StoreUpdater<NodeType extends Node = Node, EdgeType extends Edge
     reset,
     setDefaultNodesAndEdges,
   } = useStore(selector, shallow);
-  const store = useStoreApi();
+  const store = useStoreApi<NodeType, EdgeType>();
 
   useEffect(() => {
     setDefaultNodesAndEdges(props.defaultNodes, props.defaultEdges);

packages/react/src/hooks/useStore.ts

@@ -1,9 +1,10 @@
 import { useContext, useMemo } from 'react';
-import { useStoreWithEqualityFn as useZustandStore } from 'zustand/traditional';
+import { UseBoundStoreWithEqualityFn, useStoreWithEqualityFn as useZustandStore } from 'zustand/traditional';
 import { errorMessages } from '@xyflow/system';
 
 import StoreContext from '../contexts/RFStoreContext';
-import type { ReactFlowState } from '../types';
+import type { Edge, Node, ReactFlowState } from '../types';
+import { StoreApi } from 'zustand';
 
 const zustandErrorMessage = errorMessages['error001']();
 
@@ -32,8 +33,10 @@ function useStore<StateSlice = unknown>(
   return useZustandStore(store, selector, equalityFn);
 }
 
-function useStoreApi() {
+function useStoreApi<NodeType extends Node = Node, EdgeType extends Edge = Edge>() {
-  const store = useContext(StoreContext);
+  const store = useContext(StoreContext) as UseBoundStoreWithEqualityFn<
+    StoreApi<ReactFlowState<NodeType, EdgeType>>
+  > | null;
 
   if (store === null) {
     throw new Error(zustandErrorMessage);