From f2e6461634853b045a93c70a0a6ddd006677f69d Mon Sep 17 00:00:00 2001
From: Grant Steffen <gsteffen1203@gmail.com>
Date: Wed, 10 Jul 2024 13:26:25 -0500
Subject: [PATCH] Add string type input to fitNodes

---
 docs/Advanced/Refs.mdx       |  2 +-
 src/layout/useLayout.ts      | 66 ++++++++++++++++++------------------
 stories/Controls.stories.tsx |  4 +--
 3 files changed, 36 insertions(+), 36 deletions(-)

diff --git a/docs/Advanced/Refs.mdx b/docs/Advanced/Refs.mdx
index 9d55eb9..ad83dcb 100644
--- a/docs/Advanced/Refs.mdx
+++ b/docs/Advanced/Refs.mdx
@@ -81,7 +81,7 @@ export interface CanvasRef {
   /**
    * Fit a group of nodes to the viewport.
    */
-  fitNodes?: (nodeIds: string[], animated?: boolean) => void;
+  fitNodes?: (nodeIds: string | string[], animated?: boolean) => void;
 
   /**
    * Scroll to X/Y
diff --git a/src/layout/useLayout.ts b/src/layout/useLayout.ts
index da1e3ae..a6be95a 100644
--- a/src/layout/useLayout.ts
+++ b/src/layout/useLayout.ts
@@ -84,7 +84,7 @@ export interface LayoutResult {
   /**
    * Fit a group of nodes to the viewport.
    */
-  fitNodes?: (nodeIds: string[], animated?: boolean) => void;
+  fitNodes?: (nodeIds: string | string[], animated?: boolean) => void;
 
   /**
    * Scroll to X/Y
@@ -138,21 +138,21 @@ export const useLayout = ({ maxWidth, maxHeight, nodes = [], edges = [], fit, pa
         const centerX = (canvasWidth - layout.width * zoom) / 2;
         const centerY = (canvasHeight - layout.height * zoom) / 2;
         switch (position) {
-        case CanvasPosition.CENTER:
-          setXY([centerX, centerY]);
-          break;
-        case CanvasPosition.TOP:
-          setXY([centerX, 0]);
-          break;
-        case CanvasPosition.LEFT:
-          setXY([0, centerY]);
-          break;
-        case CanvasPosition.RIGHT:
-          setXY([canvasWidth - layout.width * zoom, centerY]);
-          break;
-        case CanvasPosition.BOTTOM:
-          setXY([centerX, canvasHeight - layout.height * zoom]);
-          break;
+          case CanvasPosition.CENTER:
+            setXY([centerX, centerY]);
+            break;
+          case CanvasPosition.TOP:
+            setXY([centerX, 0]);
+            break;
+          case CanvasPosition.LEFT:
+            setXY([0, centerY]);
+            break;
+          case CanvasPosition.RIGHT:
+            setXY([canvasWidth - layout.width * zoom, centerY]);
+            break;
+          case CanvasPosition.BOTTOM:
+            setXY([centerX, canvasHeight - layout.height * zoom]);
+            break;
         }
       }
     },
@@ -165,21 +165,21 @@ export const useLayout = ({ maxWidth, maxHeight, nodes = [], edges = [], fit, pa
       const scrollCenterY = (canvasHeight - height) / 2;
       if (pannable) {
         switch (position) {
-        case CanvasPosition.CENTER:
-          scrollToXY([scrollCenterX, scrollCenterY], animated);
-          break;
-        case CanvasPosition.TOP:
-          scrollToXY([scrollCenterX, 0], animated);
-          break;
-        case CanvasPosition.LEFT:
-          scrollToXY([0, scrollCenterY], animated);
-          break;
-        case CanvasPosition.RIGHT:
-          scrollToXY([canvasWidth - width, scrollCenterY], animated);
-          break;
-        case CanvasPosition.BOTTOM:
-          scrollToXY([scrollCenterX, canvasHeight - height], animated);
-          break;
+          case CanvasPosition.CENTER:
+            scrollToXY([scrollCenterX, scrollCenterY], animated);
+            break;
+          case CanvasPosition.TOP:
+            scrollToXY([scrollCenterX, 0], animated);
+            break;
+          case CanvasPosition.LEFT:
+            scrollToXY([0, scrollCenterY], animated);
+            break;
+          case CanvasPosition.RIGHT:
+            scrollToXY([canvasWidth - width, scrollCenterY], animated);
+            break;
+          case CanvasPosition.BOTTOM:
+            scrollToXY([scrollCenterX, canvasHeight - height], animated);
+            break;
         }
       }
     },
@@ -217,9 +217,9 @@ export const useLayout = ({ maxWidth, maxHeight, nodes = [], edges = [], fit, pa
    * This centers the chart on the canvas, zooms in to fit the specified nodes, and scrolls to center the nodes in the viewport
    */
   const fitNodes = useCallback(
-    (nodeIds: string[], animated = true) => {
+    (nodeIds: string | string[], animated = true) => {
       if (layout && layout.children) {
-        const nodes = nodeIds.map((nodeId) => findNode(layout.children, nodeId));
+        const nodes = Array.isArray(nodeIds) ? nodeIds.map((nodeId) => findNode(layout.children, nodeId)) : [findNode(layout.children, nodeIds)];
 
         if (nodes) {
           // center the chart
diff --git a/stories/Controls.stories.tsx b/stories/Controls.stories.tsx
index 59d0faa..f8fd3a8 100644
--- a/stories/Controls.stories.tsx
+++ b/stories/Controls.stories.tsx
@@ -223,8 +223,8 @@ export const Zoom = () => {
         <button style={{ display: 'block', width: '100%', margin: '5px 0' }} onClick={() => ref.current.zoomIn()}>Zoom In</button>
         <button style={{ display: 'block', width: '100%', margin: '5px 0' }} onClick={() => ref.current.zoomOut()}>Zoom Out</button>
         <button style={{ display: 'block', width: '100%' }} onClick={() => ref.current.fitCanvas(true)}>Fit</button>
-        <button style={{ display: 'block', width: '100%' }} onClick={() => ref.current?.fitNodes(['1'])}>Fit to Node 1</button>
-        <button style={{ display: 'block', width: '100%' }} onClick={() => ref.current?.fitNodes(['2'])}>Fit to Node 2</button>
+        <button style={{ display: 'block', width: '100%' }} onClick={() => ref.current?.fitNodes('1')}>Fit to Node 1</button>
+        <button style={{ display: 'block', width: '100%' }} onClick={() => ref.current?.fitNodes('2')}>Fit to Node 2</button>
       </pre>
       <Canvas
         maxZoom={10}