✨ Results tree uses `View` component and displays filter input fields.

1 files changed, 121 insertions, 3 deletions

diff --git a/tidal_dl_ng/helper/gui.py b/tidal_dl_ng/helper/gui.py
index 610be80..834d9b6 100644
--- a/tidal_dl_ng/helper/gui.py
+++ b/tidal_dl_ng/helper/gui.py
@@ -1,8 +1,12 @@
-from PySide6 import QtCore, QtWidgets
+import re
+
+from PySide6 import QtCore, QtGui, QtWidgets
 from tidalapi import Album, Mix, Playlist, Track, UserPlaylist, Video
 from tidalapi.artist import Artist
 from tidalapi.media import Quality
 
+from tidal_dl_ng.helper.gui import HumanProxyModel
+
 
 def get_table_data(
     item: QtWidgets.QTreeWidgetItem, column: int
@@ -18,8 +22,12 @@ def get_table_text(item: QtWidgets.QTreeWidgetItem, column: int) -> str:
     return result
 
 
-def get_results_media_item(item: QtWidgets.QTreeWidgetItem) -> Track | Video | Album | Artist | Playlist | Mix:
-    result: Track | Video | Album | Artist | Playlist | Mix = get_table_data(item, 1)
+def get_results_media_item(
+    index: QtCore.QModelIndex, proxy: HumanProxyModel, model: QtGui.QStandardItemModel
+) -> Track | Video | Album | Artist | Playlist | Mix:
+    # Switch column to "obj" column and map proxy data to our model.
+    item: QtGui.QStandardItem = model.itemFromIndex(proxy.mapToSource(index.siblingAtColumn(1)))
+    result: Track | Video | Album | Artist = item.data(QtCore.Qt.ItemDataRole.UserRole)
 
     return result
 
@@ -66,3 +74,113 @@ def set_queue_download_media(
     item: QtWidgets.QTreeWidgetItem, media: Mix | Playlist | UserPlaylist | Track | Video | Album | Artist
 ):
     set_table_data(item, media, 1)
+
+
+class FilterHeader(QtWidgets.QHeaderView):
+    filter_activated = QtCore.Signal()
+
+    def __init__(self, parent):
+        super().__init__(QtCore.Qt.Horizontal, parent)
+        self._editors = []
+        self._padding = 4
+        self.setStretchLastSection(True)
+        self.setSectionResizeMode(QtWidgets.QHeaderView.Stretch)
+        self.setDefaultAlignment(QtCore.Qt.AlignLeft | QtCore.Qt.AlignVCenter)
+        self.setSortIndicatorShown(False)
+        self.setSectionsMovable(True)
+        self.sectionResized.connect(self.adjust_positions)
+        parent.horizontalScrollBar().valueChanged.connect(self.adjust_positions)
+
+    def set_filter_boxes(self, count):
+        while self._editors:
+            editor = self._editors.pop()
+            editor.deleteLater()
+
+        for _ in range(count):
+            editor = QtWidgets.QLineEdit(self.parent())
+            editor.setPlaceholderText("Filter")
+            editor.setClearButtonEnabled(True)
+            editor.returnPressed.connect(self.filter_activated.emit)
+            self._editors.append(editor)
+
+        self.adjust_positions()
+
+    def sizeHint(self):
+        size = super().sizeHint()
+        if self._editors:
+            height = self._editors[0].sizeHint().height()
+            size.setHeight(size.height() + height + self._padding)
+
+        return size
+
+    def updateGeometries(self):
+        if self._editors:
+            height = self._editors[0].sizeHint().height()
+            self.setViewportMargins(0, 0, 0, height + self._padding)
+        else:
+            self.setViewportMargins(0, 0, 0, 0)
+
+        super().updateGeometries()
+        self.adjust_positions()
+
+    def adjust_positions(self):
+        for index, editor in enumerate(self._editors):
+            height = editor.sizeHint().height()
+
+            editor.move(self.sectionPosition(index) - self.offset() + 2, height + (self._padding // 2))
+            editor.resize(self.sectionSize(index), height)
+
+    def filter_text(self, index) -> str:
+        if 0 <= index < len(self._editors):
+            return self._editors[index].text()
+
+        return ""
+
+    def set_filter_text(self, index, text):
+        if 0 <= index < len(self._editors):
+            self._editors[index].setText(text)
+
+    def clear_filters(self):
+        for editor in self._editors:
+            editor.clear()
+
+
+class HumanProxyModel(QtCore.QSortFilterProxyModel):
+    def _human_key(self, key):
+        parts = re.split(r"(\d*\.\d+|\d+)", key)
+
+        return tuple((e.swapcase() if i % 2 == 0 else float(e)) for i, e in enumerate(parts))
+
+    def lessThan(self, source_left, source_right):
+        data_left = source_left.data()
+        data_right = source_right.data()
+
+        if isinstance(data_left, str) and isinstance(data_right, str):
+            return self._human_key(data_left) < self._human_key(data_right)
+
+        return super().lessThan(source_left, source_right)
+
+    @property
+    def filters(self):
+        if not hasattr(self, "_filters"):
+            self._filters = []
+
+        return self._filters
+
+    @filters.setter
+    def filters(self, filters):
+        print("filters() called.")
+
+        self._filters = filters
+
+        self.invalidateFilter()
+
+    def filterAcceptsRow(self, sourceRow, sourceParent) -> bool:
+        for i, text in self.filters:
+            if 0 <= i < self.columnCount():
+                ix = self.sourceModel().index(sourceRow, i, sourceParent)
+                data = ix.data()
+                if str(data) in text:
+                    return False
+
+        return True