From 1077069d7c6b352d67a6225915170d72bf2009cb Mon Sep 17 00:00:00 2001
From: "Alex \"mcmonkey\" Goodwin"
 <4000772+mcmonkey4eva@users.noreply.github.com>
Date: Tue, 24 Sep 2024 16:37:30 +0900
Subject: [PATCH] better badges for empty/loading model library folders (#953)

* better badges for empty/loading model library folders

for #945

* fix total count on loaded nodes

* fix test break

* additional test fix

* use a null prop fallback instead of having to explicitly calc

* patch
---
 src/components/common/TreeExplorer.vue        |  8 +++---
 .../common/TreeExplorerTreeNode.vue           |  2 +-
 .../__tests__/TreeExplorerTreeNode.spec.ts    |  2 +-
 .../sidebar/tabs/ModelLibrarySidebarTab.vue   | 16 ++++++++++++
 src/types/treeExplorerTypes.ts                | 25 +++++++++++--------
 5 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/src/components/common/TreeExplorer.vue b/src/components/common/TreeExplorer.vue
index 893c3f80..fb150ffb 100644
--- a/src/components/common/TreeExplorer.vue
+++ b/src/components/common/TreeExplorer.vue
@@ -82,14 +82,16 @@ const getTreeNodeIcon = (node: TreeExplorerNode) => {
 }
 const fillNodeInfo = (node: TreeExplorerNode): RenderedTreeExplorerNode => {
   const children = node.children?.map(fillNodeInfo)
+  const totalLeaves = node.leaf
+    ? 1
+    : children.reduce((acc, child) => acc + child.totalLeaves, 0)
   return {
     ...node,
     icon: getTreeNodeIcon(node),
     children,
     type: node.leaf ? 'node' : 'folder',
-    totalLeaves: node.leaf
-      ? 1
-      : children.reduce((acc, child) => acc + child.totalLeaves, 0)
+    totalLeaves,
+    badgeText: node.getBadgeText ? node.getBadgeText(node) : null
   }
 }
 const onNodeContentClick = async (
diff --git a/src/components/common/TreeExplorerTreeNode.vue b/src/components/common/TreeExplorerTreeNode.vue
index 82e5ba04..deb372fe 100644
--- a/src/components/common/TreeExplorerTreeNode.vue
+++ b/src/components/common/TreeExplorerTreeNode.vue
@@ -22,7 +22,7 @@
       </span>
       <Badge
         v-if="!props.node.leaf"
-        :value="props.node.totalLeaves"
+        :value="props.node.badgeText ?? props.node.totalLeaves"
         severity="secondary"
         class="leaf-count-badge"
       />
diff --git a/src/components/common/__tests__/TreeExplorerTreeNode.spec.ts b/src/components/common/__tests__/TreeExplorerTreeNode.spec.ts
index 3d6f181e..8a04d382 100644
--- a/src/components/common/__tests__/TreeExplorerTreeNode.spec.ts
+++ b/src/components/common/__tests__/TreeExplorerTreeNode.spec.ts
@@ -56,7 +56,7 @@ describe('TreeExplorerTreeNode', () => {
     expect(wrapper.findComponent(EditableText).props('modelValue')).toBe(
       'Test Node'
     )
-    expect(wrapper.findComponent(Badge).props()['value']).toBe(3)
+    expect(wrapper.findComponent(Badge).props()['value'].toString()).toBe('3')
   })
 
   it('makes node label editable when renamingEditingNode matches', async () => {
diff --git a/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue b/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
index a0a24d3a..c2060e87 100644
--- a/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
+++ b/src/components/sidebar/tabs/ModelLibrarySidebarTab.vue
@@ -137,6 +137,22 @@ const renderedRoot = computed<TreeExplorerNode<ComfyModelDef>>(() => {
           return 'pi pi-file'
         }
       },
+      getBadgeText: (node: TreeExplorerNode<ComfyModelDef>) => {
+        if (node.leaf) {
+          return null
+        }
+        if (node.children?.length === 1) {
+          const onlyChild = node.children[0]
+          if (onlyChild.data?.is_fake_object) {
+            if (onlyChild.data.name === '(No Content)') {
+              return '0'
+            } else if (onlyChild.data.name === 'Loading') {
+              return '?'
+            }
+          }
+        }
+        return null
+      },
       children,
       draggable: node.leaf,
       handleClick: (
diff --git a/src/types/treeExplorerTypes.ts b/src/types/treeExplorerTypes.ts
index 24b7967a..d90f5199 100644
--- a/src/types/treeExplorerTypes.ts
+++ b/src/types/treeExplorerTypes.ts
@@ -7,36 +7,39 @@ export interface TreeExplorerNode<T = any> {
   data?: T
   children?: TreeExplorerNode<T>[]
   icon?: string
+  /** Function to override what icon to use for the node */
   getIcon?: (node: TreeExplorerNode<T>) => string
-  // Function to handle renaming the node
+  /** Function to override what text to use for the leaf-count badge on a folder node */
+  getBadgeText?: (node: TreeExplorerNode<T>) => string
+  /** Function to handle renaming the node */
   handleRename?: (
     node: TreeExplorerNode<T>,
     newName: string
   ) => void | Promise<void>
-  // Function to handle deleting the node
+  /** Function to handle deleting the node */
   handleDelete?: (node: TreeExplorerNode<T>) => void | Promise<void>
-  // Function to handle adding a child node
+  /** Function to handle adding a child node */
   handleAddChild?: (
     node: TreeExplorerNode<T>,
     child: TreeExplorerNode<T>
   ) => void | Promise<void>
-  // Whether the node is draggable
+  /** Whether the node is draggable */
   draggable?: boolean
-  // Whether the node is droppable
+  /** Whether the node is droppable */
   droppable?: boolean
-  // Function to handle dropping a node
+  /** Function to handle dropping a node */
   handleDrop?: (
     node: TreeExplorerNode<T>,
     data: TreeExplorerDragAndDropData
   ) => void | Promise<void>
-  // Function to handle clicking a node
+  /** Function to handle clicking a node */
   handleClick?: (
     node: TreeExplorerNode<T>,
     event: MouseEvent
   ) => void | Promise<void>
-  // Function to handle errors
+  /** Function to handle errors */
   handleError?: (error: Error) => void
-  // Extra context menu items
+  /** Extra context menu items */
   contextMenuItems?:
     | MenuItem[]
     | ((targetNode: RenderedTreeExplorerNode) => MenuItem[])
@@ -46,8 +49,10 @@ export interface RenderedTreeExplorerNode<T = any> extends TreeExplorerNode<T> {
   children?: RenderedTreeExplorerNode<T>[]
   icon: string
   type: 'folder' | 'node'
-  // Total number of leaves in the subtree
+  /** Total number of leaves in the subtree */
   totalLeaves: number
+  /** Text to display on the leaf-count badge */
+  badgeText?: string
 }
 
 export type TreeExplorerDragAndDropData<T = any> = {