mew/ui/qt/widgets/explorertable.py

from PyQt5.QtCore import pyqtSignal
from PyQt5.QtGui import QIcon
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)

    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)
        self.trackSelected.emit(None)

    def update_tracks(self, tracks):
        # for track in tracks:
        #     pass
        self.populate()

    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 filterBySearchText(self, searchText):
        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 not in s

                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)

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

        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

                hide = tag_id and item.data(2, 0) not in track_ids

                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 tag_id:
            self.expandAll()
        else:
            self.collapseAll()

        self.setUpdatesEnabled(True)