4 files changed, 47 insertions, 25 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()){
diff --git a/beetview.h b/beetview.h
index 763dc8e..f4ae39f 100644
--- a/beetview.h
+++ b/beetview.h
@@ -6,10 +6,17 @@
 class BeetView : public QTreeView {
     Q_OBJECT
     public:
+        enum ExpandOrCollapseMode { Expand, Collapse };
         explicit BeetView(QWidget *parent = 0);
 
+    public slots:
+        void expandOrCollapse(int mode);
+
     protected:
         virtual void contextMenuEvent(QContextMenuEvent *e);
+
+    private:
+        void expandOrCollapseRecursive(const QModelIndex &idx, int mode);
 };
 
 #endif // BEETVIEW_H
diff --git a/playerwidget.cpp b/playerwidget.cpp
index e092b3b..fb8d563 100644
--- a/playerwidget.cpp
+++ b/playerwidget.cpp
@@ -164,6 +164,17 @@ void PlayerWidget::setupGui(QSplashScreen *splash){
         QAction *addToPlayListAndClearA = new QAction(QIcon(":/belly_right_and_clear.png"), tr("Clear and add"), this);
         connect(addToPlayListAndClearA, &QAction::triggered, this, &PlayerWidget::addToPlayListAndClear);
         curW->view()->addAction(addToPlayListAndClearA);
+        QStandardItemModel *model = curW->model();
+        QModelIndex rootIdx = model->invisibleRootItem()->index();
+        if(model->hasChildren(model->index(0, 0, rootIdx))){
+            QAction *expandA = new QAction(Helper::iconFromQChar(QChar(0x2640), 90), tr("Expand"), this);
+            connect(expandA, &QAction::triggered, [curW] { curW->view()->expandOrCollapse(BeetView::Expand);} );
+            curW->view()->addAction(Helper::createSeparator(this));
+            curW->view()->addAction(expandA);
+            QAction *collapseA = new QAction(Helper::iconFromQChar(QChar(0x2642), 120), tr("Collapse"), this);
+            connect(collapseA, &QAction::triggered, [curW] { curW->view()->expandOrCollapse(BeetView::Collapse);} );
+            curW->view()->addAction(collapseA);
+        }
     }
 
     //left widget
@@ -326,8 +337,6 @@ void PlayerWidget::createActions(){
     connect(addToPlayListA, &QAction::triggered, this, &PlayerWidget::addToPlayList);
     QAction *addToPlayListAndClearA = new QAction(QIcon(":/belly_right_and_clear.png"), tr("Clear and add"), this);
     connect(addToPlayListAndClearA, &QAction::triggered, this, &PlayerWidget::addToPlayListAndClear);
-    QAction *expandA = new QAction(Helper::iconFromQChar(QChar(0x2640), 90), tr("Expand"), this);
-    connect(expandA, &QAction::triggered, this, &PlayerWidget::expand);
     QAction *removeFromPlayListA = new QAction(QIcon(":/belly_left.png"), tr("Remove from playlist"), this);
     connect(removeFromPlayListA, &QAction::triggered, this, &PlayerWidget::removeFromPlayList);
     QAction *clearPlayListA = new QAction(QIcon(":/delete.png"), tr("Clear Playlist"), this);
@@ -960,27 +969,6 @@ void PlayerWidget::continuePlaying(QMediaPlayer::MediaStatus state){
     }
 }
 
-void PlayerWidget::expand(){
-    CollectionWidget *curWidget = qobject_cast<CollectionWidget*>(mCollectionStack->currentWidget());
-    QModelIndexList sel = curWidget->view()->selectionModel()->selectedRows();
-    for(const QModelIndex &i : sel){
-        curWidget->view()->expand(i);
-        expandRecursive(i);
-    }
-}
-
-void PlayerWidget::expandRecursive(const QModelIndex &idx){
-    const QStandardItemModel *model = static_cast<const QStandardItemModel*>(idx.model());
-    CollectionWidget *curWidget = qobject_cast<CollectionWidget*>(mCollectionStack->currentWidget());
-    QStandardItem *item = model->itemFromIndex(idx);
-    for(int i = 0; i < item->rowCount(); ++i){
-        QModelIndex cur = model->indexFromItem(item->child(i, 0));
-        if(cur.isValid()){
-            curWidget->view()->expand(cur);
-        }
-    }
-}
-
 void PlayerWidget::readSettings(){
     mStarting = true;
     QSettings s;
diff --git a/playerwidget.h b/playerwidget.h
index 60e0df9..84b2ba3 100644
--- a/playerwidget.h
+++ b/playerwidget.h
@@ -70,7 +70,6 @@ class PlayerWidget : public QWidget {
         void setPosition(qint64 pos);
         void slide(int value);
         void continuePlaying(QMediaPlayer::MediaStatus state);
-        void expand();
         void readSettings();
         void writeSettings();
         void aboutDlg();
@@ -93,7 +92,6 @@ class PlayerWidget : public QWidget {
         void play(const QString &fullPath);
         void playUrl(const QString &url);
         void advance(int numSongs);
-        void expandRecursive(const QModelIndex &idx);
         void adjustVolume(int by);
         void fillWithText(QTextEdit *te, const TagLib::FileRef &fr);
         void setNowPlaying(const QString &what);