From 4c35893e0f45744a484753b4345c09771282fd75 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christopher=20M=C3=B6ller?= <chrtze@gmail.com>
Date: Tue, 6 Dec 2022 18:08:51 +0100
Subject: [PATCH] feat(node-toolbar): make nodeId optional and get id from
 context

---
 examples/vite-app/src/examples/NodeToolbar/CustomNode.tsx | 2 +-
 packages/node-toolbar/src/NodeToolbar.tsx                 | 7 +++++--
 packages/node-toolbar/src/types.ts                        | 2 +-
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/examples/vite-app/src/examples/NodeToolbar/CustomNode.tsx b/examples/vite-app/src/examples/NodeToolbar/CustomNode.tsx
index 4f8ab270..ca2ac0b7 100644
--- a/examples/vite-app/src/examples/NodeToolbar/CustomNode.tsx
+++ b/examples/vite-app/src/examples/NodeToolbar/CustomNode.tsx
@@ -4,7 +4,7 @@ import { Handle, Position, NodeProps, NodeToolbar } from 'reactflow';
 const CustomNode: FC<NodeProps> = ({ id, data }) => {
   return (
     <>
-      <NodeToolbar nodeId={id} isVisible={data.toolbarVisible} position={data.toolbarPosition}>
+      <NodeToolbar isVisible={data.toolbarVisible} position={data.toolbarPosition}>
         <button>delete</button>
         <button>copy</button>
         <button>expand</button>
diff --git a/packages/node-toolbar/src/NodeToolbar.tsx b/packages/node-toolbar/src/NodeToolbar.tsx
index 1dfd7669..5c6cbc86 100644
--- a/packages/node-toolbar/src/NodeToolbar.tsx
+++ b/packages/node-toolbar/src/NodeToolbar.tsx
@@ -8,6 +8,7 @@ import {
   Rect,
   Position,
   internalsSymbol,
+  useNodeId,
 } from '@reactflow/core';
 import cc from 'classcat';
 import shallow from 'zustand/shallow';
@@ -72,9 +73,11 @@ function NodeToolbar({
   offset = 10,
   ...rest
 }: NodeToolbarProps) {
+  const contextNodeId = useNodeId();
+
   const nodesSelector = useCallback(
     (state: ReactFlowState): Node[] => {
-      const nodeIds: string[] = typeof nodeId === 'string' ? [nodeId] : nodeId;
+      const nodeIds = Array.isArray(nodeId) ? nodeId : [nodeId || contextNodeId || ''];
 
       return nodeIds.reduce<Node[]>((acc, id) => {
         const node = state.nodeInternals.get(id);
@@ -84,7 +87,7 @@ function NodeToolbar({
         return acc;
       }, [] as Node[]);
     },
-    [nodeId]
+    [nodeId, contextNodeId]
   );
   const nodes = useStore(nodesSelector, nodesEqualityFn);
   const { transform, nodeOrigin, selectedNodesCount } = useStore(storeSelector, shallow);
diff --git a/packages/node-toolbar/src/types.ts b/packages/node-toolbar/src/types.ts
index a9f547ad..e2b0d73c 100644
--- a/packages/node-toolbar/src/types.ts
+++ b/packages/node-toolbar/src/types.ts
@@ -2,7 +2,7 @@ import { Position } from '@reactflow/core';
 import type { HTMLAttributes } from 'react';
 
 export type NodeToolbarProps = HTMLAttributes<HTMLDivElement> & {
-  nodeId: string | string[];
+  nodeId?: string | string[];
   isVisible?: boolean;
   position?: Position;
   offset?: number;