feat(react): infer node types from passed nodes

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

@@ -17,9 +17,9 @@ import { useDrag } from '../../hooks/useDrag';
 import { useUpdateNodePositions } from '../../hooks/useUpdateNodePositions';
 import { handleNodeClick } from '../Nodes/utils';
 import { arrowKeyDiffs, builtinNodeTypes } from './utils';
-import type { NodeWrapperProps } from '../../types';
+import type { Node, NodeWrapperProps } from '../../types';
 
-export function NodeWrapper({
+export function NodeWrapper<NodeType extends Node>({
   id,
   onClick,
   onMouseEnter,
@@ -40,9 +40,9 @@ export function NodeWrapper({
   nodeExtent,
   nodeOrigin,
   onError,
-}: NodeWrapperProps) {
+}: NodeWrapperProps<NodeType>) {
   const { node, positionAbsoluteX, positionAbsoluteY, zIndex, isParent } = useStore((s) => {
-    const node = s.nodeLookup.get(id)!;
+    const node = s.nodeLookup.get(id)! as NodeType;
 
     const positionAbsolute = nodeExtent
       ? clampPosition(node.computed?.positionAbsolute, nodeExtent)

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

@@ -14,8 +14,8 @@ import { useUpdateNodePositions } from '../../hooks/useUpdateNodePositions';
 import { arrowKeyDiffs } from '../NodeWrapper/utils';
 import type { Node, ReactFlowState } from '../../types';
 
-export type NodesSelectionProps = {
-  onSelectionContextMenu?: (event: MouseEvent, nodes: Node[]) => void;
+export type NodesSelectionProps<NodeType> = {
+  onSelectionContextMenu?: (event: MouseEvent, nodes: NodeType[]) => void;
   noPanClassName?: string;
   disableKeyboardA11y: boolean;
 };
@@ -32,7 +32,11 @@ const selector = (s: ReactFlowState) => {
   };
 };
 
-export function NodesSelection({ onSelectionContextMenu, noPanClassName, disableKeyboardA11y }: NodesSelectionProps) {
+export function NodesSelection<NodeType extends Node>({
+  onSelectionContextMenu,
+  noPanClassName,
+  disableKeyboardA11y,
+}: NodesSelectionProps<NodeType>) {
   const store = useStoreApi();
   const { width, height, transformString, userSelectionActive } = useStore(selector, shallow);
   const updatePositions = useUpdateNodePositions();
@@ -58,7 +62,7 @@ export function NodesSelection({ onSelectionContextMenu, noPanClassName, disable
   const onContextMenu = onSelectionContextMenu
     ? (event: MouseEvent) => {
         const selectedNodes = store.getState().nodes.filter((n) => n.selected);
-        onSelectionContextMenu(event, selectedNodes);
+        onSelectionContextMenu(event, selectedNodes as NodeType[]);
       }
     : undefined;
 

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

@@ -68,7 +68,9 @@ const reactFlowFieldsToTrack = [
 ] as const;
 
 type ReactFlowFieldsToTrack = (typeof reactFlowFieldsToTrack)[number];
-type StoreUpdaterProps = Pick<ReactFlowProps, ReactFlowFieldsToTrack> & { rfId: string };
+type StoreUpdaterProps<NodeType extends Node = Node> = Pick<ReactFlowProps<NodeType>, ReactFlowFieldsToTrack> & {
+  rfId: string;
+};
 
 // rfId doesn't exist in ReactFlowProps, but it's one of the fields we want to update
 const fieldsToTrack = [...reactFlowFieldsToTrack, 'rfId'] as const;
@@ -84,7 +86,7 @@ const selector = (s: ReactFlowState) => ({
   reset: s.reset,
 });
 
-export function StoreUpdater(props: StoreUpdaterProps) {
+export function StoreUpdater<NodeType extends Node = Node>(props: StoreUpdaterProps<NodeType>) {
   const {
     setNodes,
     setEdges,
@@ -106,7 +108,7 @@ export function StoreUpdater(props: StoreUpdaterProps) {
     };
   }, []);
 
-  const previousFields = useRef<Partial<StoreUpdaterProps>>({
+  const previousFields = useRef<Partial<StoreUpdaterProps<NodeType>>>({
     // these are values that are also passed directly to other components
     // than the StoreUpdater. We can reduce the number of setStore calls
     // by setting the same values here as prev fields.