From 07ddd97650283064cf7b706cfb94fdf68693c4fe Mon Sep 17 00:00:00 2001
From: moklick <info@moritzklack.com>
Date: Mon, 19 Feb 2024 10:45:48 +0100
Subject: [PATCH] refactor(react): add generics to useStoreApi

---
 .../src/additional-components/MiniMap/MiniMap.tsx     |  4 ++--
 .../react/src/components/NodesSelection/index.tsx     |  5 ++---
 packages/react/src/components/StoreUpdater/index.tsx  |  2 +-
 packages/react/src/hooks/useStore.ts                  | 11 +++++++----
 4 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/packages/react/src/additional-components/MiniMap/MiniMap.tsx b/packages/react/src/additional-components/MiniMap/MiniMap.tsx
index c4bbf13b..17a3f7f2 100644
--- a/packages/react/src/additional-components/MiniMap/MiniMap.tsx
+++ b/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;
diff --git a/packages/react/src/components/NodesSelection/index.tsx b/packages/react/src/components/NodesSelection/index.tsx
index c0423827..6c037ba5 100644
--- a/packages/react/src/components/NodesSelection/index.tsx
+++ b/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;
 
diff --git a/packages/react/src/components/StoreUpdater/index.tsx b/packages/react/src/components/StoreUpdater/index.tsx
index 85159721..7917e3d6 100644
--- a/packages/react/src/components/StoreUpdater/index.tsx
+++ b/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);
diff --git a/packages/react/src/hooks/useStore.ts b/packages/react/src/hooks/useStore.ts
index 9a3a5e6b..8a600706 100644
--- a/packages/react/src/hooks/useStore.ts
+++ b/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() {
-  const store = useContext(StoreContext);
+function useStoreApi<NodeType extends Node = Node, EdgeType extends Edge = Edge>() {
+  const store = useContext(StoreContext) as UseBoundStoreWithEqualityFn<
+    StoreApi<ReactFlowState<NodeType, EdgeType>>
+  > | null;
 
   if (store === null) {
     throw new Error(zustandErrorMessage);