mew/ui/qt/widgets/explorertable.py

from PyQt5.QtCore import pyqtSignal
from PyQt5.QtGui import QIcon, QBrush, QColor
from PyQt5.QtWidgets import QTreeWidget, QTreeWidgetItem, QTreeWidgetItemIterator

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

logger = Logger.get()


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

    def __init__(self, parent):
        super(QTreeWidget, self).__init__(parent)
        self.setColumnCount(3)
        self.setColumnWidth(0, 54)
        self.hideColumn(2)
        self.itemDoubleClicked.connect(self._itemDoubleClicked)
        self.artist_items, self.album_items, self.track_items = {}, {}, {}

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

        self.artist_items, self.album_items, self.track_items = {}, {}, {}

        tracks = track_repo.get_all()

        self.setUpdatesEnabled(False)
        self.update_tracks(tracks)
        self.setUpdatesEnabled(True)
        self.trackSelected.emit(None)

    def update_tracks(self, tracks):

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

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

        for track in tracks:

            if track.id not in self.track_items:
                if track.artist and track.album:
                    parent = self.album_items[(track.artist, track.album)]
                elif track.artist:
                    parent = self.artist_items[track.artist]
                else:
                    parent = self

                item = QTreeWidgetItem(parent)
                item.setIcon(0, QIcon(":/img/rsc/music.png"))
                item.setData(2, 0, track.id)
                self.track_items[track.id] = item

            if track.status != track.STATUS_UNAVAILABLE:
                title = track.title
                color = QColor(20, 20, 20)
            else:
                title = f"[Introuvable] {track.title}"
                color = QColor(200, 20, 20)

            self.track_items[track.id].setText(1, title)
            self.track_items[track.id].setForeground(0, QBrush(color))
            self.track_items[track.id].setForeground(1, QBrush(color))

    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, _):
        try:
            track_id = selected.indexes()[2].data(2)
        except IndexError:
            return
        if track_id:
            track = TrackRepository().get_by_id(int(track_id))
        else:
            track = None
        self.trackSelected.emit(track)

    def _itemDoubleClicked(self, item, _):
        track_id = item.data(2, 0)
        if track_id:
            track = TrackRepository().get_by_id(int(track_id))
            self.trackDoubleClicked.emit(track)

    def filterBy(self, tag_id=None, searchText=""):

        track_ids = []
        if tag_id:
            track_repo = TrackRepository()
            track_ids = [t.id for t in track_repo.get_by_tag_id(tag_id)]

        searchText = norm_search(searchText)

        self.setUpdatesEnabled(False)

        iterator = QTreeWidgetItemIterator(self)
        while True:
            item = iterator.value()
            if item is not None:

                if not item.data(2, 0):
                    # not a track
                    item.setHidden(True)
                    iterator += 1
                    continue

                track_item = item
                parent_item = item.parent()
                grand_parent_item = parent_item.parent() if parent_item else None

                keys = [track_item.text(1)]
                if parent_item:
                    keys.append(parent_item.text(1))
                if grand_parent_item:
                    keys.append(grand_parent_item.text(1))

                s = norm_search(" ".join(keys))
                hide = (searchText and searchText not in s) \
                       or (tag_id and item.data(2, 0) not in track_ids)

                track_item.setHidden(hide)

                iterator += 1
            else:
                break

        iterator = QTreeWidgetItemIterator(self)
        while True:
            item = iterator.value()
            if item is not None:
                if not item.data(2, 0) or item.isHidden():
                    # not a track or hidden
                    iterator += 1
                    continue

                if item.parent():
                    item.parent().setHidden(False)
                if item.parent() and item.parent().parent():
                    item.parent().parent().setHidden(False)

                iterator += 1
            else:
                break

        if searchText:
            self.expandAll()
        else:
            self.collapseAll()

        self.setUpdatesEnabled(True)