1 files changed, 29 insertions, 0 deletions
diff --git a/beetview.cpp b/beetview.cpp
index 875eb5a..82dd86b 100644
--- a/beetview.cpp
+++ b/beetview.cpp
@@ -1,11 +1,40 @@
 #include <QMenu>
 #include <QContextMenuEvent>
+#include <QStandardItemModel>
+#include <QSortFilterProxyModel>
 
 #include "beetview.h"
 
 BeetView::BeetView(QWidget *parent) : QTreeView(parent) {
 }
 
+void BeetView::expandOrCollapse(int mode){
+    QModelIndexList sel = selectionModel()->selectedRows();
+    for(QModelIndex i : sel){
+        if(mode == Expand){
+            expand(i);
+        }else if(mode == Collapse){
+            collapse(i);
+        }
+        expandOrCollapseRecursive(i, mode);
+    }
+}
+
+void BeetView::expandOrCollapseRecursive(const QModelIndex &idx, int mode){
+    const QSortFilterProxyModel *pm = qobject_cast<const QSortFilterProxyModel*>(model());
+    for(int i = 0; i < pm->rowCount(idx); ++i){
+        QModelIndex curIdx = pm->index(i, 0, idx);
+        if(curIdx.isValid()){
+            if(mode == Expand){
+                expand(curIdx);
+            }else if(mode == Collapse){
+                collapse(curIdx);
+            }
+            expandOrCollapseRecursive(curIdx, mode);
+        }
+    }
+}
+
 void BeetView::contextMenuEvent(QContextMenuEvent *e){
     QMenu ctxMenu;
     for(QAction *a : actions()){