Column search fixes

Author: mrvisscher

activity_browser/ui/core/tree_model.py

@@ -1,11 +1,12 @@
 from typing import Optional
-from collections import defaultdict
 
-from loguru import logger
 import pandas as pd
-from PySide6.QtCore import QModelIndex, Qt
+
+from PySide6 import QtGui
+from PySide6.QtCore import QModelIndex, Qt, QAbstractItemModel
 from PySide6.QtWidgets import QWidget
-from PySide6.QtCore import QAbstractItemModel
+
+from activity_browser.ui.icons import qicons
 
 
 class TreeNode:
@@ -54,6 +55,8 @@ def __init__(self, df: pd.DataFrame = None, parent: Optional[QWidget] = None, ch
         super().__init__(parent)
         self.df = df if df is not None else pd.DataFrame()
         self.df_query: dict[str, str] = {"model": "index == index"}  # dictionary where queries can be registered
+        self.filtered_columns: set[int] = set()  # set of column indices that have active filters, only used for the header icon
+
         self.lazy = chunk_size > 0
         self.chunk_size = chunk_size
 
@@ -269,13 +272,22 @@ def isBranchNode(self, index: QModelIndex) -> bool:
         return not node.is_leaf
 
     def headerData(self, section: int, orientation: Qt.Orientation = Qt.Horizontal, role: int = Qt.DisplayRole):
-        if orientation == Qt.Vertical or not role == Qt.DisplayRole:
+        if orientation == Qt.Vertical:
             return None
-        
-        if section == 0:
-            return ""
-              
-        return self.df.columns[section - 1]
+
+        if role == Qt.DisplayRole:
+            if section == 0:
+                return ""
+
+            return self.df.columns[section - 1]
+
+        if role == Qt.ItemDataRole.FontRole and section in self.filtered_columns:
+            font = QtGui.QFont()
+            font.setUnderline(True)
+            return font
+
+        if role == Qt.ItemDataRole.DecorationRole and section in self.filtered_columns:
+            return qicons.filter
 
     def canFetchMore(self, parent: QModelIndex) -> bool:
         """Check if this parent has more children that can be loaded."""
@@ -333,7 +345,8 @@ def build_node_hierarchy(self, pandas_index: pd.Index) -> None:
 
         # Create a mapping from full path to DataFrame position
         path_to_position = {}
-        for df_pos, row_tuple in enumerate(idx_df.itertuples(index=False, name=None)):
+        for row_tuple in idx_df.itertuples(index=False, name=None):
+            df_pos = self.df.index.get_loc(row_tuple)
             path_to_position[row_tuple] = df_pos
 
         # Process each level to build the hierarchy

activity_browser/ui/widgets/tree_view.py

@@ -1,12 +1,10 @@
 from loguru import logger
 
-import pandas as pd
-
 from qtpy import QtWidgets, QtCore, QtGui
-from qtpy.QtCore import Qt
 
 from activity_browser.ui import delegates, core
-from .item_model import ABItemModel
+
+from .line_edit import ABLineEdit
 
 
 
@@ -26,11 +24,11 @@ def __init__(self, pos: QtCore.QPoint, view: "ABNewTreeView"):
             col_index = view.columnAt(pos.x())
             col_name = model.columns()[col_index]
 
-            search_box = QtWidgets.QLineEdit(self)
+            search_box = ABLineEdit(self)
             search_box.setText(view.columnFilters.get(col_name, ""))
             search_box.setPlaceholderText("Search")
             search_box.selectAll()
-            search_box.textChanged.connect(lambda query: view.setColumnFilter(col_name, query))
+            search_box.textChangedDebounce.connect(lambda query: view.setColumnFilter(col_name, query))
             widget_action = QtWidgets.QWidgetAction(self)
             widget_action.setDefaultWidget(search_box)
             self.addAction(widget_action)
@@ -103,7 +101,7 @@ def setModel(self, model):
         self.updateIndexColumnVisibility()
         self.updateBranchSpanning()
 
-    def model(self) -> ABItemModel:
+    def model(self) -> core.ABTreeModel:
         return super().model()
 
     # === Functionality related to contextmenus
@@ -122,10 +120,10 @@ def setColumnFilter(self, column_name: str, query: str):
 
         if query:
             self.columnFilters[column_name] = query
-            # self.model().filtered_columns.add(col_index)
+            self.model().filtered_columns.add(col_index)
         elif column_name in self.columnFilters:
             del self.columnFilters[column_name]
-            # self.model().filtered_columns.discard(col_index)
+            self.model().filtered_columns.discard(col_index)
 
         self.applyFilter()
 
@@ -144,7 +142,7 @@ def buildQuery(self) -> str:
                 del self.columnFilters[col]
 
         for col, query in self.columnFilters.items():
-            q = f"({col}.astype('str').str.contains('{self.format_query(query)}'))"
+            q = f"({col}.astype('str').str.contains('{self.format_query(query)}', False))"
             queries.append(q)
 
         # query for the all filter