Bind extra context menu items on TreeExplorerNode interface (#886)

Comfy-Org · Sep 19, 2024 · a57c958 · a57c958
1 parent b6dbe8f
commit a57c958
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 40 deletions.
diff --git a/src/components/common/TreeExplorer.vue b/src/components/common/TreeExplorer.vue
@@ -52,9 +52,6 @@ provide('selectionKeys', selectionKeys)
 const props = defineProps<{
   roots: TreeExplorerNode[]
   class?: string
-  extraMenuItems?:
-    | MenuItem[]
-    | ((targetNode: RenderedTreeExplorerNode) => MenuItem[])
 }>()
 const emit = defineEmits<{
   (e: 'nodeClick', node: RenderedTreeExplorerNode, event: MouseEvent): void
@@ -101,6 +98,13 @@ const onNodeContentClick = (e: MouseEvent, node: RenderedTreeExplorerNode) => {
 const menu = ref(null)
 const menuTargetNode = ref<RenderedTreeExplorerNode | null>(null)
 provide('menuTargetNode', menuTargetNode)
+const extraMenuItems = computed(() => {
+  return menuTargetNode.value?.contextMenuItems
+    ? typeof menuTargetNode.value.contextMenuItems === 'function'
+      ? menuTargetNode.value.contextMenuItems(menuTargetNode.value)
+      : menuTargetNode.value.contextMenuItems
+    : []
+})
 const renameEditingNode = ref<RenderedTreeExplorerNode | null>(null)
 provide('renameEditingNode', renameEditingNode)
 
@@ -126,11 +130,7 @@ const menuItems = computed<MenuItem[]>(() =>
       command: () => deleteCommand(menuTargetNode.value),
       visible: menuTargetNode.value?.handleDelete !== undefined
     },
-    ...(props.extraMenuItems
-      ? typeof props.extraMenuItems === 'function'
-        ? props.extraMenuItems(menuTargetNode.value)
-        : props.extraMenuItems
-      : [])
+    ...extraMenuItems.value
   ].map((menuItem) => ({
     ...menuItem,
     command: wrapCommandWithErrorHandler(menuItem.command)

diff --git a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
@@ -49,7 +49,6 @@
             class="node-lib-tree-explorer mt-1"
             :roots="renderedRoot.children"
             v-model:expandedKeys="expandedKeys"
-            @nodeClick="handleNodeClick"
           >
             <template #node="{ node }">
               <NodeTreeLeaf :node="node" />

diff --git a/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue b/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue
@@ -4,7 +4,6 @@
     ref="treeExplorerRef"
     :roots="renderedBookmarkedRoot.children"
     :expandedKeys="expandedKeys"
-    :extraMenuItems="extraMenuItems"
   >
     <template #folder="{ node }">
       <NodeTreeFolder :node="node" />
@@ -89,6 +88,37 @@ watch(
     }
   }
 )
+
+const { t } = useI18n()
+const extraMenuItems = (
+  menuTargetNode: RenderedTreeExplorerNode<ComfyNodeDefImpl>
+) => [
+  {
+    label: t('newFolder'),
+    icon: 'pi pi-folder-plus',
+    command: () => {
+      addNewBookmarkFolder(menuTargetNode)
+    },
+    visible: !menuTargetNode?.leaf
+  },
+  {
+    label: t('customize'),
+    icon: 'pi pi-palette',
+    command: () => {
+      const customization =
+        nodeBookmarkStore.bookmarksCustomization[menuTargetNode.data.nodePath]
+      initialIcon.value =
+        customization?.icon || nodeBookmarkStore.defaultBookmarkIcon
+      initialColor.value =
+        customization?.color || nodeBookmarkStore.defaultBookmarkColor
+
+      showCustomizationDialog.value = true
+      customizationTargetNodePath.value = menuTargetNode.data.nodePath
+    },
+    visible: !menuTargetNode?.leaf
+  }
+]
+
 const renderedBookmarkedRoot = computed<TreeExplorerNode<ComfyNodeDefImpl>>(
   () => {
     const fillNodeInfo = (
@@ -145,6 +175,7 @@ const renderedBookmarkedRoot = computed<TreeExplorerNode<ComfyNodeDefImpl>>(
             toggleNodeOnEvent(e, node)
           }
         },
+        contextMenuItems: extraMenuItems,
         ...(node.leaf
           ? {}
           : {
@@ -201,34 +232,4 @@ const updateCustomization = (icon: string, color: string) => {
     )
   }
 }
-
-const { t } = useI18n()
-const extraMenuItems = computed(
-  () => (menuTargetNode: RenderedTreeExplorerNode<ComfyNodeDefImpl>) => [
-    {
-      label: t('newFolder'),
-      icon: 'pi pi-folder-plus',
-      command: () => {
-        addNewBookmarkFolder(menuTargetNode)
-      },
-      visible: !menuTargetNode?.leaf
-    },
-    {
-      label: t('customize'),
-      icon: 'pi pi-palette',
-      command: () => {
-        const customization =
-          nodeBookmarkStore.bookmarksCustomization[menuTargetNode.data.nodePath]
-        initialIcon.value =
-          customization?.icon || nodeBookmarkStore.defaultBookmarkIcon
-        initialColor.value =
-          customization?.color || nodeBookmarkStore.defaultBookmarkColor
-
-        showCustomizationDialog.value = true
-        customizationTargetNodePath.value = menuTargetNode.data.nodePath
-      },
-      visible: !menuTargetNode?.leaf
-    }
-  ]
-)
 </script>
diff --git a/src/types/treeExplorerTypes.ts b/src/types/treeExplorerTypes.ts
@@ -1,3 +1,5 @@
+import type { MenuItem } from 'primevue/menuitem'
+
 export interface TreeExplorerNode<T = any> {
   key: string
   label: string
@@ -28,6 +30,10 @@ export interface TreeExplorerNode<T = any> {
   handleClick?: (node: TreeExplorerNode<T>, event: MouseEvent) => void
   // Function to handle errors
   handleError?: (error: Error) => void
+  // Extra context menu items
+  contextMenuItems?:
+    | MenuItem[]
+    | ((targetNode: RenderedTreeExplorerNode) => MenuItem[])
 }
 
 export interface RenderedTreeExplorerNode<T = any> extends TreeExplorerNode<T> {