From c1f30b96f9bdd8a71029dbc45675cf816d760e6c Mon Sep 17 00:00:00 2001
From: christian-byrne <abolkonsky.rem@gmail.com>
Date: Tue, 14 Jan 2025 18:19:15 -0700
Subject: [PATCH] Add key function param in node tree build

---
 .../sidebar/tabs/NodeLibrarySidebarTab.vue    | 30 ++++++++++++++++++-
 src/locales/en/main.json                      |  3 +-
 src/stores/nodeDefStore.ts                    | 23 +++++++++++---
 3 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
index 7b1385604f..152302e205 100644
--- a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
+++ b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
@@ -20,6 +20,14 @@
         @click="alphabeticalSort = !alphabeticalSort"
         v-tooltip.bottom="$t('sideToolbar.nodeLibraryTab.sortOrder')"
       />
+      <Button
+        class="grouping-button"
+        :icon="groupBySource ? 'pi pi-list' : 'pi pi-list-check'"
+        text
+        severity="secondary"
+        @click="groupBySource = !groupBySource"
+        v-tooltip.bottom="$t('sideToolbar.nodeLibraryTab.groupingType')"
+      />
     </template>
     <template #header>
       <SearchBox
@@ -66,7 +74,7 @@ import Button from 'primevue/button'
 import Divider from 'primevue/divider'
 import Popover from 'primevue/popover'
 import type { TreeNode } from 'primevue/treenode'
-import { Ref, computed, nextTick, ref } from 'vue'
+import { Ref, computed, nextTick, ref, watch } from 'vue'
 
 import SearchBox from '@/components/common/SearchBox.vue'
 import { SearchFilter } from '@/components/common/SearchFilterChip.vue'
@@ -101,6 +109,21 @@ const nodeBookmarkTreeExplorerRef = ref<InstanceType<
 > | null>(null)
 const searchFilter = ref(null)
 const alphabeticalSort = ref(false)
+const groupBySource = ref(false)
+
+const createSourceKey = (nodeDef: ComfyNodeDefImpl) => {
+  const sourcePath = nodeDef.python_module.split('.')
+  const pathWithoutCategory = nodeDef.nodePath.split('/').slice(1)
+  return [...sourcePath, ...pathWithoutCategory]
+}
+
+watch(groupBySource, (newValue) => {
+  if (newValue) {
+    nodeDefStore.setKeyFunction(createSourceKey)
+  } else {
+    nodeDefStore.setKeyFunction(null)
+  }
+})
 
 const searchQuery = ref<string>('')
 
@@ -194,4 +217,9 @@ const onRemoveFilter = (filterAndValue) => {
   }
   handleSearch(searchQuery.value)
 }
+
+// This can be added if the persistent state is not desirable:
+// onBeforeUnmount(() => {
+// nodeDefStore.setKeyFunction(null)
+// })
 </script>
diff --git a/src/locales/en/main.json b/src/locales/en/main.json
index 8fd22e1e31..854b4a33d5 100644
--- a/src/locales/en/main.json
+++ b/src/locales/en/main.json
@@ -234,7 +234,8 @@
     "openWorkflow": "Open workflow in local file system",
     "newBlankWorkflow": "Create a new blank workflow",
     "nodeLibraryTab": {
-      "sortOrder": "Sort Order"
+      "sortOrder": "Sort Order",
+      "groupingType": "Grouping Type"
     },
     "modelLibrary": "Model Library",
     "downloads": "Downloads",
diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts
index f6acc73226..60dbc44544 100644
--- a/src/stores/nodeDefStore.ts
+++ b/src/stores/nodeDefStore.ts
@@ -305,9 +305,13 @@ export const SYSTEM_NODE_DEFS: Record<string, ComfyNodeDef> = {
   }
 }
 
-export function buildNodeDefTree(nodeDefs: ComfyNodeDefImpl[]): TreeNode {
-  return buildTree(nodeDefs, (nodeDef: ComfyNodeDefImpl) =>
+export function buildNodeDefTree(
+  nodeDefs: ComfyNodeDefImpl[],
+  keyFunction: (nodeDef: ComfyNodeDefImpl) => string[] = (nodeDef) =>
     nodeDef.nodePath.split('/')
+): TreeNode {
+  return buildTree(nodeDefs, (nodeDef: ComfyNodeDefImpl) =>
+    keyFunction(nodeDef)
   )
 }
 
@@ -326,12 +330,16 @@ export function createDummyFolderNodeDef(folderPath: string): ComfyNodeDefImpl {
   } as ComfyNodeDef)
 }
 
+const getCategoryKeys = (nodeDef: ComfyNodeDefImpl) =>
+  nodeDef.nodePath.split('/')
+
 export const useNodeDefStore = defineStore('nodeDef', () => {
   const nodeDefsByName = ref<Record<string, ComfyNodeDefImpl>>({})
   const nodeDefsByDisplayName = ref<Record<string, ComfyNodeDefImpl>>({})
   const showDeprecated = ref(false)
   const showExperimental = ref(false)
 
+  const keyFunction = ref(getCategoryKeys)
   const nodeDefs = computed(() => Object.values(nodeDefsByName.value))
   const nodeDataTypes = computed(() => {
     const types = new Set<string>()
@@ -355,7 +363,13 @@ export const useNodeDefStore = defineStore('nodeDef', () => {
   const nodeSearchService = computed(
     () => new NodeSearchService(visibleNodeDefs.value)
   )
-  const nodeTree = computed(() => buildNodeDefTree(visibleNodeDefs.value))
+  const nodeTree = computed(() =>
+    buildNodeDefTree(visibleNodeDefs.value, keyFunction.value)
+  )
+
+  function setKeyFunction(callback: (nodeDef: ComfyNodeDefImpl) => string[]) {
+    keyFunction.value = callback ?? getCategoryKeys
+  }
 
   function updateNodeDefs(nodeDefs: ComfyNodeDef[]) {
     const newNodeDefsByName: Record<string, ComfyNodeDefImpl> = {}
@@ -398,7 +412,8 @@ export const useNodeDefStore = defineStore('nodeDef', () => {
 
     updateNodeDefs,
     addNodeDef,
-    fromLGraphNode
+    fromLGraphNode,
+    setKeyFunction
   }
 })