mew/ui/qt/widgets/explorertable.py

from PyQt5.QtCore import pyqtSlot, pyqtSignal
from PyQt5.QtGui import QIcon
from PyQt5.QtWidgets import QTreeWidget, QTreeWidgetItem

from core import db
from core.logging_ import Logger
from core.repositories import TrackRepository

logger = Logger.get()


class ExplorerTable(QTreeWidget):
    trackSelected = pyqtSignal(object)

    def __init__(self, parent):
        super(QTreeWidget, self).__init__(parent)
        self.setColumnCount(3)
        self.setColumnWidth(0, 54)
        self.hideColumn(2)

    def populate(self):
        session = db.Session()
        track_repo = TrackRepository(session)
        self.clear()

        artist_folders = {}
        album_folders = {}
        tracks = track_repo.get_all()

        self.setUpdatesEnabled(False)

        for track in tracks:
            if track.artist and track.artist not in artist_folders:
                item = QTreeWidgetItem(self)
                item.setIcon(0, QIcon(":/img/rsc/artist.png"))
                item.setText(1, track.artist)
                artist_folders[track.artist] = item

        for track in tracks:
            # list albums with artists
            if track.artist and track.album and (track.artist, track.album) not in album_folders:
                item = QTreeWidgetItem(artist_folders[track.artist])
                item.setIcon(0, QIcon(":/img/rsc/album.png"))
                item.setText(1, track.album)
                album_folders[(track.artist, track.album)] = item

        for track in tracks:
            if track.artist and track.album:
                parent = album_folders[(track.artist, track.album)]
            elif track.artist:
                parent = artist_folders[track.artist]
            else:
                parent = self

            item = QTreeWidgetItem(parent)
            item.setIcon(0, QIcon(":/img/rsc/music.png"))
            item.setText(1, track.title)
            item.setData(2, 0, track.id)

        self.setUpdatesEnabled(True)

    def selected_track(self):
        track_id = self.selectionModel().selection().indexes()[2].data(2)
        if not track_id:
            return None
        return TrackRepository().get_by_id(int(track_id))

    def selectionChanged(self, selected, _):
        track_id = selected.indexes()[2].data(2)
        if track_id:
            track = TrackRepository().get_by_id(int(track_id))
        else:
            track = None
        self.trackSelected.emit(track)