add expand/collapse buttons to directories in tree view

javascript/clusterize.js

@@ -188,22 +188,26 @@ class Clusterize {
     }
 
     async setMaxItems(max_items) {
+        /** Sets the new max number of items.
+         *
+         *  This is used to control the scroll bar's length.
+         *
+         *  Returns whether the number of max items changed.
+         */
         if (!this.setup_has_run || !this.enabled) {
             this.#max_items = max_items;
-            return;
+            return this.#max_items !== max_items;
         }
-
-        if (max_items === this.#max_items) {
+        if (this.#max_items === max_items) {
             // No change. do nothing.
-            return;
+            return false;
         }
-
         // If the number of items changed, we need to update the cluster.
         this.#max_items = max_items;
         await this.refresh();
-
         // Apply sort to the updated data.
         await this.sortData();
+        return true;
     }
 
     // ==== PRIVATE FUNCTIONS ====

javascript/extraNetworks.js

@@ -938,15 +938,14 @@ function extraNetworksTreeDirectoryOnClick(event, btn, tabname_full) {
 
     const tab = extra_networks_tabs[tabname_full];
 
-    if (true_targ.matches(".tree-list-item-action--leading, .tree-list-item-action-chevron")) {
+    const prev_selected_elem = gradioApp().querySelector(".tree-list-item[data-selected='']");
+    if (true_targ.matches(".tree-list-item-action--leading .tree-list-item-action-chevron")) {
         // If user clicks on the chevron, then we do not select the folder.
-        const prev_selected_elem = gradioApp().querySelector(".tree-list-item[data-selected='']");
         tab.tree_list.onRowExpandClick(div_id, btn);
-        const selected_elem = gradioApp().querySelector(".tree-list-item[data-selected='']");
-        if (isElement(prev_selected_elem) && !isElement(selected_elem)) {
-            // is a selected element was removed, clear filter.
-            tab.updateSearch("");
-        }
+    } else if (true_targ.matches(".tree-list-item-action--trailing .tree-list-item-action-expand")) {
+        tab.tree_list.onExpandAllClick(div_id);
+    } else if (true_targ.matches(".tree-list-item-action--trailing .tree-list-item-action-collapse")) {
+        tab.tree_list.onCollapseAllClick(div_id);
     } else {
         // user clicked anywhere else on the row
         tab.tree_list.onRowSelected(div_id, btn);
@@ -963,6 +962,11 @@ function extraNetworksTreeDirectoryOnClick(event, btn, tabname_full) {
         }
         tab.updateSearch("selected" in btn.dataset ? btn.dataset.path : "");
     }
+    const selected_elem = gradioApp().querySelector(".tree-list-item[data-selected='']");
+    if (isElement(prev_selected_elem) && !isElement(selected_elem)) {
+        // if a selected element was removed, clear filter.
+        tab.updateSearch("");
+    }
 }
 
 function extraNetworksTreeOnClick(event, tabname_full) {

javascript/extraNetworksClusterize.js

@@ -338,6 +338,58 @@ class ExtraNetworksClusterizeTreeList extends ExtraNetworksClusterize {
         return max_width;
     }
 
+    async onExpandAllClick(div_id) {
+        if (!keyExistsLogError(this.data_obj, div_id)) {
+            return;
+        }
+
+        const _expand = (parent_id) => {
+            const this_obj = this.data_obj[parent_id];
+            this_obj.visible = true;
+            this_obj.expanded = true;
+            for (const child_id of this_obj.children) {
+                _expand(child_id);
+            }
+        };
+
+        this.data_obj[div_id].expanded = true;
+        for (const child_id of this.data_obj[div_id].children) {
+            _expand(child_id);
+        }
+
+        const new_len = Object.values(this.data_obj).filter(v => v.visible).length;
+        const max_items_changed = await this.setMaxItems(new_len);
+        if (!max_items_changed) {
+            await this.refresh(true);
+        }
+    }
+
+    async onCollapseAllClick(div_id) {
+        if (!keyExistsLogError(this.data_obj, div_id)) {
+            return;
+        }
+
+        const _collapse = (parent_id) => {
+            const this_obj = this.data_obj[parent_id];
+            this_obj.visible = false;
+            this_obj.expanded = false;
+            for (const child_id of this_obj.children) {
+                _collapse(child_id);
+            }
+        };
+
+        this.data_obj[div_id].expanded = false;
+        for (const child_id of this.data_obj[div_id].children) {
+            _collapse(child_id);
+        }
+
+        const new_len = Object.values(this.data_obj).filter(v => v.visible).length;
+        const max_items_changed = await this.setMaxItems(new_len);
+        if (!max_items_changed) {
+            await this.refresh(true);
+        }
+    }
+
     async onRowExpandClick(div_id, elem) {
         /** Expands or collapses a row to show/hide children. */
         if (!keyExistsLogError(this.data_obj, div_id)) {
@@ -353,7 +405,10 @@ class ExtraNetworksClusterizeTreeList extends ExtraNetworksClusterize {
         }
 
         const new_len = Object.values(this.data_obj).filter(v => v.visible).length;
-        await this.setMaxItems(new_len);
+        const max_items_changed = await this.setMaxItems(new_len);
+        if (!max_items_changed) {
+            await this.refresh(true);
+        }
     }
 
     async initData() {