feat(hooks): add useReactFlow hook, refactor types and properties

src/hooks/useNodeBounds.ts

@@ -1,31 +0,0 @@
-import { useCallback } from 'react';
-import shallow from 'zustand/shallow';
-
-import { NodeBounds } from '../types';
-import { useStore } from '../store';
-
-function useNodeBounds(id: string): NodeBounds | null {
-  const nodeBounds = useStore(
-    useCallback(
-      (s) => {
-        const nodeItem = s.nodeInternals.get(id);
-
-        if (!nodeItem) {
-          return null;
-        }
-
-        return {
-          ...nodeItem.positionAbsolute,
-          width: nodeItem.width ?? null,
-          height: nodeItem.height ?? null,
-        };
-      },
-      [id]
-    ),
-    shallow
-  );
-
-  return nodeBounds;
-}
-
-export default useNodeBounds;

src/hooks/useOnPaneReadyHandler.ts

@@ -1,67 +1,19 @@
 import { useEffect, useRef } from 'react';
 
-import { pointToRendererPoint } from '../utils/graph';
-import { useStoreApi } from '../store';
-import useZoomPanHelper from '../hooks/useZoomPanHelper';
-import { OnPaneReady, XYPosition, Node, Edge, FlowExportObject } from '../types';
+import useReactFlow from './useReactFlow';
+import { OnPaneReady } from '../types';
 
 function useOnPaneReadyHandler(onPaneReady: OnPaneReady<any> | undefined) {
+  const ReactFlowInstance = useReactFlow();
   const isInitialized = useRef<boolean>(false);
-  const store = useStoreApi();
-  const { zoomIn, zoomOut, zoomTo, getZoom, setTransform, getTransform, setCenter, fitView, initialized } =
-    useZoomPanHelper();
 
   useEffect(() => {
-    if (!isInitialized.current && initialized) {
-      if (onPaneReady) {
-        const project = (position: XYPosition): XYPosition => {
-          const { transform, snapToGrid, snapGrid } = store.getState();
-          return pointToRendererPoint(position, transform, snapToGrid, snapGrid);
-        };
-
-        const getNodes = (): Node[] => {
-          const { nodeInternals } = store.getState();
-          // @TODO: work with nodeInternals instead of converting it to an array
-          const nodes = Array.from(nodeInternals).map(([_, node]) => node);
-          return nodes.map((n) => ({ ...n }));
-        };
-
-        const getEdges = (): Edge[] => {
-          const { edges = [] } = store.getState();
-          return edges.map((e) => ({ ...e }));
-        };
-
-        const toObject = (): FlowExportObject => {
-          const { nodeInternals, edges = [], transform } = store.getState();
-          // @TODO: work with nodeInternals instead of converting it to an array
-          const nodes = Array.from(nodeInternals).map(([_, node]) => node);
-          return {
-            nodes: nodes.map((n) => ({ ...n })),
-            edges: edges.map((e) => ({ ...e })),
-            position: [transform[0], transform[1]],
-            zoom: transform[2],
-          };
-        };
-
-        onPaneReady({
-          fitView: (params = { padding: 0.1 }) => fitView(params),
-          zoomIn,
-          zoomOut,
-          zoomTo,
-          getZoom,
-          setTransform,
-          getTransform,
-          setCenter,
-          project,
-          getNodes,
-          getEdges,
-          toObject,
-        });
-      }
+    if (!isInitialized.current && ReactFlowInstance.viewportInitialized && onPaneReady) {
+      onPaneReady(ReactFlowInstance);
 
       isInitialized.current = true;
     }
-  }, [onPaneReady, zoomIn, zoomOut, zoomTo, setTransform, fitView, initialized]);
+  }, [onPaneReady, ReactFlowInstance]);
 }
 
 export default useOnPaneReadyHandler;

src/hooks/useReactFlow.ts

@@ -0,0 +1,86 @@
+import { useCallback } from 'react';
+import useViewportHelper from './useViewportHelper';
+import { useStoreApi } from '../store';
+import { ReactFlowInstance, Instance } from '../types';
+
+export default function useReactFlow<NodeData, EdgeData>(): ReactFlowInstance<NodeData, EdgeData> {
+  const { initialized: viewportInitialized, ...viewportHelperFunctions } = useViewportHelper();
+  const store = useStoreApi();
+
+  const getNodes = useCallback<Instance.GetNodes<NodeData>>(() => {
+    const { nodeInternals } = store.getState();
+    const nodes = Array.from(nodeInternals.values());
+    return nodes.map((n) => ({ ...n }));
+  }, []);
+
+  const getNode = useCallback<Instance.GetNode<NodeData>>((id) => {
+    const { nodeInternals } = store.getState();
+    return nodeInternals.get(id);
+  }, []);
+
+  const getEdges = useCallback<Instance.GetEdges<EdgeData>>(() => {
+    const { edges = [] } = store.getState();
+    return edges.map((e) => ({ ...e }));
+  }, []);
+
+  const getEdge = useCallback<Instance.GetEdge<EdgeData>>((id) => {
+    const { edges = [] } = store.getState();
+    return edges.find((e) => e.id === id);
+  }, []);
+
+  const setNodes = useCallback<Instance.SetNodes<NodeData>>((payload) => {
+    const { nodeInternals, setNodes } = store.getState();
+    const nodes = Array.from(nodeInternals.values());
+    const nextNodes = typeof payload === 'function' ? payload(nodes) : payload;
+    setNodes(nextNodes);
+  }, []);
+
+  const setEdges = useCallback<Instance.SetEdges<EdgeData>>((payload) => {
+    const { edges = [], setEdges } = store.getState();
+    const nextEdges = typeof payload === 'function' ? payload(edges) : payload;
+    setEdges(nextEdges);
+  }, []);
+
+  const addNodes = useCallback<Instance.AddNodes<NodeData>>((payload) => {
+    const nodes = Array.isArray(payload) ? payload : [payload];
+    const { nodeInternals, setNodes } = store.getState();
+    const currentNodes = Array.from(nodeInternals.values());
+    const nextNodes = [...currentNodes, ...nodes];
+    setNodes(nextNodes);
+  }, []);
+
+  const addEdges = useCallback<Instance.AddEdges<EdgeData>>((payload) => {
+    const nextEdges = Array.isArray(payload) ? payload : [payload];
+    const { edges = [], setEdges } = store.getState();
+    setEdges([...edges, ...nextEdges]);
+  }, []);
+
+  const toObject = useCallback<Instance.ToObject<NodeData, EdgeData>>(() => {
+    const { nodeInternals, edges = [], transform } = store.getState();
+    const nodes = Array.from(nodeInternals.values());
+    const [x, y, zoom] = transform;
+    return {
+      nodes: nodes.map((n) => ({ ...n })),
+      edges: edges.map((e) => ({ ...e })),
+      viewport: {
+        x,
+        y,
+        zoom,
+      },
+    };
+  }, []);
+
+  return {
+    ...viewportHelperFunctions,
+    viewportInitialized,
+    getNodes,
+    getNode,
+    getEdges,
+    getEdge,
+    setNodes,
+    setEdges,
+    addNodes,
+    addEdges,
+    toObject,
+  };
+}

src/hooks/useUpdateEdges.ts

@@ -1,23 +0,0 @@
-import { useMemo } from 'react';
-import { useStoreApi } from '../store';
-import { Edge } from '../types';
-
-type FunctionParams<T> = (nodes: Edge<T>[], setEdges: (nodes: Edge<T>[]) => void) => void;
-
-function useUpdateEdges<T>() {
-  const store = useStoreApi();
-
-  const update = useMemo(() => {
-    return (action: FunctionParams<T>): void => {
-      const { edges, setEdges } = store.getState();
-
-      if (typeof action === 'function') {
-        action(edges, setEdges);
-      }
-    };
-  }, []);
-
-  return update;
-}
-
-export default useUpdateEdges;

src/hooks/useUpdateNodes.ts

@@ -1,24 +0,0 @@
-import { useMemo } from 'react';
-import { useStoreApi } from '../store';
-import { Node } from '../types';
-
-type FunctionParams<T> = (nodes: Node<T>[], setNodes: (nodes: Node<T>[]) => void) => void;
-
-function useUpdateNodes<T>() {
-  const store = useStoreApi();
-
-  const update = useMemo(() => {
-    return (action: FunctionParams<T>): void => {
-      const { nodeInternals, setNodes } = store.getState();
-      const nodes = Array.from(nodeInternals.values()) as Node<T>[];
-
-      if (typeof action === 'function') {
-        action(nodes, setNodes);
-      }
-    };
-  }, []);
-
-  return update;
-}
-
-export default useUpdateNodes;

src/hooks/useViewportHelper.ts

@@ -5,17 +5,17 @@ import { Selection as D3Selection } from 'd3';
 
 import { useStoreApi, useStore } from '../store';
 import { getRectOfNodeInternals, pointToRendererPoint, getTransformForBounds } from '../utils/graph';
-import { FitViewParams, FlowTransform, ZoomPanHelperFunctions, ReactFlowState, Rect, XYPosition } from '../types';
+import { FitViewParams, Viewport, ViewportHelperFunctions, ReactFlowState, Rect, XYPosition } from '../types';
 
 const DEFAULT_PADDING = 0.1;
 
-const initialZoomPanHelper: ZoomPanHelperFunctions = {
+const initialViewportHelper: ViewportHelperFunctions = {
   zoomIn: () => {},
   zoomOut: () => {},
   zoomTo: (_: number) => {},
   getZoom: () => 1,
-  setTransform: (_: FlowTransform) => {},
-  getTransform: () => ({ x: 0, y: 0, zoom: 1 }),
+  setViewport: (_: Viewport) => {},
+  getViewport: () => ({ x: 0, y: 0, zoom: 1 }),
   fitView: (_: FitViewParams = { padding: DEFAULT_PADDING, includeHiddenNodes: false }) => {},
   setCenter: (_: number, __: number) => {},
   fitBounds: (_: Rect) => {},
@@ -32,11 +32,11 @@ const getTransition = (selection: D3Selection<Element, unknown, null, undefined>
   return selection.transition().duration(duration);
 };
 
-const useZoomPanHelper = (): ZoomPanHelperFunctions => {
+const useViewportHelper = (): ViewportHelperFunctions => {
   const store = useStoreApi();
   const { d3Zoom, d3Selection } = useStore(selector, shallow);
 
-  const zoomPanHelperFunctions = useMemo<ZoomPanHelperFunctions>(() => {
+  const viewportHelperFunctions = useMemo<ViewportHelperFunctions>(() => {
     if (d3Selection && d3Zoom) {
       return {
         zoomIn: (options) => d3Zoom.scaleBy(getTransition(d3Selection, options?.duration), 1.2),
@@ -46,11 +46,11 @@ const useZoomPanHelper = (): ZoomPanHelperFunctions => {
           const [, , zoom] = store.getState().transform;
           return zoom;
         },
-        setTransform: (transform, options) => {
+        setViewport: (transform, options) => {
           const nextTransform = zoomIdentity.translate(transform.x, transform.y).scale(transform.zoom);
           d3Zoom.transform(getTransition(d3Selection, options?.duration), nextTransform);
         },
-        getTransform: () => {
+        getViewport: () => {
           const [x, y, zoom] = store.getState().transform;
           return { x, y, zoom };
         },
@@ -109,10 +109,10 @@ const useZoomPanHelper = (): ZoomPanHelperFunctions => {
       };
     }
 
-    return initialZoomPanHelper;
+    return initialViewportHelper;
   }, [d3Zoom, d3Selection]);
 
-  return zoomPanHelperFunctions;
+  return viewportHelperFunctions;
 };
 
-export default useZoomPanHelper;
+export default useViewportHelper;