From d9ae6cb395b513894d877cad0680eaba4c334b7c Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Thu, 20 Mar 2025 12:03:47 -0400
Subject: [PATCH] [TS] Use custom TreeNode type (#3164)

---
 .../dialog/content/SettingDialogContent.vue    | 18 ++++++------------
 .../sidebar/tabs/ModelLibrarySidebarTab.vue    |  4 +---
 .../sidebar/tabs/NodeLibrarySidebarTab.vue     |  3 +--
 .../sidebar/tabs/WorkflowsSidebarTab.vue       |  4 +---
 .../nodeLibrary/NodeBookmarkTreeExplorer.vue   |  3 +--
 src/composables/useTreeExpansion.ts            |  3 ++-
 src/stores/nodeBookmarkStore.ts                |  2 +-
 src/stores/nodeDefStore.ts                     |  2 +-
 src/stores/settingStore.ts                     |  2 +-
 src/types/treeExplorerTypes.ts                 |  9 ++++++---
 src/utils/treeUtil.ts                          |  4 ++--
 tests-ui/tests/utils/treeUtilTest.test.ts      | 18 ++----------------
 12 files changed, 25 insertions(+), 47 deletions(-)

diff --git a/src/components/dialog/content/SettingDialogContent.vue b/src/components/dialog/content/SettingDialogContent.vue
index 7cfab04549..75ded5b9f2 100644
--- a/src/components/dialog/content/SettingDialogContent.vue
+++ b/src/components/dialog/content/SettingDialogContent.vue
@@ -157,9 +157,7 @@ const categories = computed<SettingTreeNode[]>(() =>
   ].map((node) => ({
     ...node,
     translatedLabel: t(
-      // @ts-expect-error fixme ts strict error
       `settingsCategories.${normalizeI18nKey(node.label)}`,
-      // @ts-expect-error fixme ts strict error
       node.label
     )
   }))
@@ -177,16 +175,12 @@ onMounted(() => {
 })
 
 const sortedGroups = (category: SettingTreeNode): ISettingGroup[] => {
-  // @ts-expect-error fixme ts strict error
-  return (
-    [...(category.children ?? [])]
-      // @ts-expect-error fixme ts strict error
-      .sort((a, b) => a.label.localeCompare(b.label))
-      .map((group) => ({
-        label: group.label,
-        settings: flattenTree<SettingParams>(group)
-      }))
-  )
+  return [...(category.children ?? [])]
+    .sort((a, b) => a.label.localeCompare(b.label))
+    .map((group) => ({
+      label: group.label,
+      settings: flattenTree<SettingParams>(group)
+    }))
 }
 
 const searchQuery = ref<string>('')
diff --git a/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue b/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
index 2ead1a5188..915b930eb6 100644
--- a/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
+++ b/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
@@ -46,7 +46,6 @@
 
 <script setup lang="ts">
 import Button from 'primevue/button'
-import type { TreeNode } from 'primevue/treenode'
 import { computed, nextTick, onMounted, ref, toRef, watch } from 'vue'
 
 import SearchBox from '@/components/common/SearchBox.vue'
@@ -64,6 +63,7 @@ import {
 } from '@/stores/modelStore'
 import { useModelToNodeStore } from '@/stores/modelToNodeStore'
 import { useSettingStore } from '@/stores/settingStore'
+import type { TreeNode } from '@/types/treeExplorerTypes'
 import type { TreeExplorerNode } from '@/types/treeExplorerTypes'
 import { isElectron } from '@/utils/envUtil'
 import { buildTree } from '@/utils/treeUtil'
@@ -116,13 +116,11 @@ const renderedRoot = computed<TreeExplorerNode<ModelOrFolder>>(() => {
 
     return {
       key: node.key,
-      // @ts-expect-error fixme ts strict error
       label: model
         ? nameFormat === 'title'
           ? model.title
           : model.simplified_file_name
         : node.label,
-      // @ts-expect-error fixme ts strict error
       leaf: node.leaf,
       data: node.data,
       getIcon() {
diff --git a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
index 0284792815..b988c47346 100644
--- a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
+++ b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
@@ -65,7 +65,6 @@
 import Button from 'primevue/button'
 import Divider from 'primevue/divider'
 import Popover from 'primevue/popover'
-import type { TreeNode } from 'primevue/treenode'
 import { Ref, computed, h, nextTick, ref, render } from 'vue'
 
 import SearchBox from '@/components/common/SearchBox.vue'
@@ -84,6 +83,7 @@ import {
   buildNodeDefTree,
   useNodeDefStore
 } from '@/stores/nodeDefStore'
+import type { TreeNode } from '@/types/treeExplorerTypes'
 import type { TreeExplorerNode } from '@/types/treeExplorerTypes'
 import { sortedTree } from '@/utils/treeUtil'
 
@@ -114,7 +114,6 @@ const renderedRoot = computed<TreeExplorerNode<ComfyNodeDefImpl>>(() => {
     return {
       key: node.key,
       label: node.leaf ? node.data.display_name : node.label,
-      // @ts-expect-error fixme ts strict error
       leaf: node.leaf,
       data: node.data,
       getIcon() {
diff --git a/src/components/sidebar/tabs/WorkflowsSidebarTab.vue b/src/components/sidebar/tabs/WorkflowsSidebarTab.vue
index d113c9aff7..90ecfebf8b 100644
--- a/src/components/sidebar/tabs/WorkflowsSidebarTab.vue
+++ b/src/components/sidebar/tabs/WorkflowsSidebarTab.vue
@@ -125,7 +125,6 @@
 <script setup lang="ts">
 import Button from 'primevue/button'
 import ConfirmDialog from 'primevue/confirmdialog'
-import type { TreeNode } from 'primevue/treenode'
 import { computed, nextTick, onMounted, ref } from 'vue'
 import { useI18n } from 'vue-i18n'
 
@@ -145,6 +144,7 @@ import {
 } from '@/stores/workflowStore'
 import { ComfyWorkflow } from '@/stores/workflowStore'
 import { useWorkspaceStore } from '@/stores/workspaceStore'
+import type { TreeNode } from '@/types/treeExplorerTypes'
 import { TreeExplorerNode } from '@/types/treeExplorerTypes'
 import { appendJsonExt } from '@/utils/formatUtil'
 import { buildTree, sortedTree } from '@/utils/treeUtil'
@@ -267,9 +267,7 @@ const renderTreeNode = (
 
   return {
     key: node.key,
-    // @ts-expect-error fixme ts strict error
     label: node.label,
-    // @ts-expect-error fixme ts strict error
     leaf: node.leaf,
     data: node.data,
     children,
diff --git a/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue b/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue
index db5ec862c2..b93a577153 100644
--- a/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue
+++ b/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue
@@ -22,7 +22,6 @@
 </template>
 
 <script setup lang="ts">
-import type { TreeNode } from 'primevue/treenode'
 import { computed, h, nextTick, ref, render, watch } from 'vue'
 import { useI18n } from 'vue-i18n'
 
@@ -35,6 +34,7 @@ import { useTreeExpansion } from '@/composables/useTreeExpansion'
 import { useLitegraphService } from '@/services/litegraphService'
 import { useNodeBookmarkStore } from '@/stores/nodeBookmarkStore'
 import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
+import type { TreeNode } from '@/types/treeExplorerTypes'
 import type {
   RenderedTreeExplorerNode,
   TreeExplorerDragAndDropData,
@@ -131,7 +131,6 @@ const renderedBookmarkedRoot = computed<TreeExplorerNode<ComfyNodeDefImpl>>(
       return {
         key: node.key,
         label: node.leaf ? node.data.display_name : node.label,
-        // @ts-expect-error fixme ts strict error
         leaf: node.leaf,
         data: node.data,
         getIcon() {
diff --git a/src/composables/useTreeExpansion.ts b/src/composables/useTreeExpansion.ts
index 9ede6654cb..16f1151c5a 100644
--- a/src/composables/useTreeExpansion.ts
+++ b/src/composables/useTreeExpansion.ts
@@ -1,6 +1,7 @@
-import type { TreeNode } from 'primevue/treenode'
 import { Ref } from 'vue'
 
+import type { TreeNode } from '@/types/treeExplorerTypes'
+
 export function useTreeExpansion(expandedKeys: Ref<Record<string, boolean>>) {
   const toggleNode = (node: TreeNode) => {
     if (node.key && typeof node.key === 'string') {
diff --git a/src/stores/nodeBookmarkStore.ts b/src/stores/nodeBookmarkStore.ts
index b54341c9a2..f629ecb428 100644
--- a/src/stores/nodeBookmarkStore.ts
+++ b/src/stores/nodeBookmarkStore.ts
@@ -1,9 +1,9 @@
 import _ from 'lodash'
 import { defineStore } from 'pinia'
-import type { TreeNode } from 'primevue/treenode'
 import { computed } from 'vue'
 
 import type { BookmarkCustomization } from '@/schemas/apiSchema'
+import type { TreeNode } from '@/types/treeExplorerTypes'
 
 import { useNodeDefStore } from './nodeDefStore'
 import { ComfyNodeDefImpl, createDummyFolderNodeDef } from './nodeDefStore'
diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts
index 87265b2bae..fc67ebb52a 100644
--- a/src/stores/nodeDefStore.ts
+++ b/src/stores/nodeDefStore.ts
@@ -1,7 +1,6 @@
 import type { LGraphNode } from '@comfyorg/litegraph'
 import axios from 'axios'
 import { defineStore } from 'pinia'
-import type { TreeNode } from 'primevue/treenode'
 import { computed, ref } from 'vue'
 
 import { transformNodeDefV1ToV2 } from '@/schemas/nodeDef/migration'
@@ -24,6 +23,7 @@ import {
   NodeSourceType,
   getNodeSource
 } from '@/types/nodeSource'
+import type { TreeNode } from '@/types/treeExplorerTypes'
 import { buildTree } from '@/utils/treeUtil'
 
 export class ComfyNodeDefImpl implements ComfyNodeDefV1, ComfyNodeDefV2 {
diff --git a/src/stores/settingStore.ts b/src/stores/settingStore.ts
index de4e810fa8..fb6f0240c5 100644
--- a/src/stores/settingStore.ts
+++ b/src/stores/settingStore.ts
@@ -1,12 +1,12 @@
 import _ from 'lodash'
 import { defineStore } from 'pinia'
-import type { TreeNode } from 'primevue/treenode'
 import { computed, ref } from 'vue'
 
 import type { Settings } from '@/schemas/apiSchema'
 import { api } from '@/scripts/api'
 import { app } from '@/scripts/app'
 import type { SettingParams } from '@/types/settingTypes'
+import type { TreeNode } from '@/types/treeExplorerTypes'
 import { buildTree } from '@/utils/treeUtil'
 
 export const getSettingInfo = (setting: SettingParams) => {
diff --git a/src/types/treeExplorerTypes.ts b/src/types/treeExplorerTypes.ts
index bce7e8a44e..9c6aaf6281 100644
--- a/src/types/treeExplorerTypes.ts
+++ b/src/types/treeExplorerTypes.ts
@@ -1,10 +1,13 @@
 import type { MenuItem } from 'primevue/menuitem'
+import type { TreeNode as PrimeVueTreeNode } from 'primevue/treenode'
 import type { InjectionKey, ModelRef } from 'vue'
 
-export interface TreeExplorerNode<T = any> {
-  key: string
+export interface TreeNode extends PrimeVueTreeNode {
   label: string
-  leaf: boolean
+  children?: TreeNode[]
+}
+
+export interface TreeExplorerNode<T = any> extends TreeNode {
   data?: T
   children?: TreeExplorerNode<T>[]
   icon?: string
diff --git a/src/utils/treeUtil.ts b/src/utils/treeUtil.ts
index 6190cd24aa..f7aaabc74c 100644
--- a/src/utils/treeUtil.ts
+++ b/src/utils/treeUtil.ts
@@ -1,4 +1,4 @@
-import type { TreeNode } from 'primevue/treenode'
+import type { TreeNode } from '@/types/treeExplorerTypes'
 
 export function buildTree<T>(items: T[], key: (item: T) => string[]): TreeNode {
   const root: TreeNode = {
@@ -45,7 +45,7 @@ export function flattenTree<T>(tree: TreeNode): T[] {
   while (stack.length) {
     const node = stack.pop()!
     if (node.leaf && node.data) result.push(node.data)
-    stack.push(...(node.children || []))
+    stack.push(...(node.children ?? []))
   }
   return result
 }
diff --git a/tests-ui/tests/utils/treeUtilTest.test.ts b/tests-ui/tests/utils/treeUtilTest.test.ts
index ef991c1e10..9c367a30d6 100644
--- a/tests-ui/tests/utils/treeUtilTest.test.ts
+++ b/tests-ui/tests/utils/treeUtilTest.test.ts
@@ -1,6 +1,6 @@
-import { TreeNode } from 'primevue/treenode'
 import { describe, expect, it } from 'vitest'
 
+import { TreeNode } from '@/types/treeExplorerTypes'
 import { buildTree, sortedTree } from '@/utils/treeUtil'
 
 describe('buildTree', () => {
@@ -84,6 +84,7 @@ describe('sortedTree', () => {
     const node: TreeNode = {
       key: 'root',
       label: 'root',
+      leaf: false,
       children: [createNode('c'), createNode('a'), createNode('b')]
     }
 
@@ -91,21 +92,6 @@ describe('sortedTree', () => {
     expect(result.children?.map((c) => c.label)).toEqual(['a', 'b', 'c'])
   })
 
-  it('should handle undefined labels', () => {
-    const node: TreeNode = {
-      key: 'root',
-      label: 'root',
-      children: [
-        { key: '1', label: 'b' },
-        { key: '2', label: 'a' },
-        { key: '3', label: undefined }
-      ]
-    }
-
-    const result = sortedTree(node)
-    expect(result.children?.map((c) => c.label)).toEqual([undefined, 'a', 'b'])
-  })
-
   describe('with groupLeaf=true', () => {
     it('should group folders before files', () => {
       const node: TreeNode = {