add filter

components/main_interface.py

@@ -1,9 +1,9 @@
 from datetime import datetime, timezone
 from zoneinfo import ZoneInfo
 from PySide6.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QAbstractItemView
-from PySide6.QtCore import QAbstractTableModel, QModelIndex, Qt
+from PySide6.QtCore import QAbstractTableModel, QModelIndex, Qt, QSortFilterProxyModel
 from qfluentwidgets import (
-    PushButton, TableView, FluentIcon as Fi,
+    PushButton, TableView, FluentIcon as Fi, PillPushButton
 )
 
 from common.utils import SAFE_MAP_ICON, SAFE_MAP
@@ -79,6 +79,22 @@ class ExtensionsDataTable(QAbstractTableModel):
         return dt.strftime(custom_format)
 
 
+class SafeFilterProxyModel(QSortFilterProxyModel):
+
+    def __init__(self, parent=None):
+        super().__init__(parent)
+        self.accepted_status = set()
+
+    def set_accepted_status(self, status: list[int]):
+        self.accepted_status = set(status)
+        self.invalidateFilter()
+
+    def filterAcceptsRow(self, source_row: int, source_parent, /):
+        index = self.sourceModel().index(source_row, 2, source_parent)
+        ext_d: dict[str, str | int] = self.sourceModel().data(index, Qt.ItemDataRole.UserRole)
+        return ext_d["SAFE"] in self.accepted_status
+
+
 class MainInterface(QWidget):
 
     def __init__(self, ext_data: list[dict[str, str | int]] = None, parent=None):
@@ -92,15 +108,33 @@ class MainInterface(QWidget):
         self.pbn_add = PushButton(Fi.ADD, "添加", self)
         self.pbn_delete = PushButton(Fi.DELETE, "删除", self)
 
+        safe_checks = [
+            ("安全", 1), ("未知", 0), ("不安全", -1), ("未记录", -2),
+        ]
+        self.safe_switches: list[PillPushButton] = []
+        for text, m in safe_checks:
+            c = PillPushButton(self)
+            c.setText(text)
+            c.setIcon(SAFE_MAP_ICON[m])
+            c.setProperty("mark", m)
+            c.toggled.connect(self.update_filter)
+            self.safe_switches.append(c)
+
         self.hly_top.addWidget(self.pbn_refresh)
         self.hly_top.addWidget(self.pbn_add)
         self.hly_top.addWidget(self.pbn_delete)
         self.hly_top.addStretch(1)
+        for c in self.safe_switches:
+            self.hly_top.addWidget(c)
         self.vly_m.addLayout(self.hly_top)
 
         self.tbv_m = TableView(self)
         self.ext_model = ExtensionsDataTable(ext_data, self)
-        self.tbv_m.setModel(self.ext_model)
+        self.fil_model = SafeFilterProxyModel(self)
+        self.fil_model.setSourceModel(self.ext_model)
+
+        self.tbv_m.setModel(self.fil_model)
+        self.tbv_m.setSortingEnabled(True)
         self.tbv_m.setEditTriggers(QAbstractItemView.EditTrigger.NoEditTriggers)
         self.tbv_m.setSelectionBehavior(QAbstractItemView.SelectionBehavior.SelectRows)  # 整行选中
         self.tbv_m.setSelectionMode(QAbstractItemView.SelectionMode.SingleSelection)  # 只允许单选
@@ -113,4 +147,16 @@ class MainInterface(QWidget):
         self.tbv_m.setColumnWidth(1, 200)
         self.tbv_m.setColumnWidth(3, 180)
 
-        self.vly_m.addWidget(self.tbv_m)
+        self.vly_m.addWidget(self.tbv_m)
+
+        # 放到最后
+        for c in self.safe_switches:
+            c.setChecked(True)
+
+    def update_filter(self, checked: bool):
+        accepted_status = []
+        for c in self.safe_switches:
+            if c.isChecked():
+                accepted_status.append(c.property("mark"))
+
+        self.fil_model.set_accepted_status(accepted_status)