From d91e6d07e2bca1d1426e0c63e1f3b1e8d18bde8c Mon Sep 17 00:00:00 2001
From: Jakub Trllo <43494761+iLLiCiTiT@users.noreply.github.com>
Date: Tue, 1 Apr 2025 12:08:49 +0200
Subject: [PATCH] added basic implementation of filtering by ids

---
 .../ayon_core/tools/utils/folders_widget.py   | 39 +++++++++++++++++--
 client/ayon_core/tools/utils/tasks_widget.py  | 36 ++++++++++++++++-
 2 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/client/ayon_core/tools/utils/folders_widget.py b/client/ayon_core/tools/utils/folders_widget.py
index b3b5cebb0e..2e2c5b5d77 100644
--- a/client/ayon_core/tools/utils/folders_widget.py
+++ b/client/ayon_core/tools/utils/folders_widget.py
@@ -1,4 +1,5 @@
 import collections
+from typing import Optional
 
 from qtpy import QtWidgets, QtGui, QtCore
 
@@ -33,7 +34,7 @@ class FoldersQtModel(QtGui.QStandardItemModel):
     refreshed = QtCore.Signal()
 
     def __init__(self, controller):
-        super(FoldersQtModel, self).__init__()
+        super().__init__()
 
         self._controller = controller
         self._items_by_id = {}
@@ -334,6 +335,29 @@ class FoldersQtModel(QtGui.QStandardItemModel):
         self.refreshed.emit()
 
 
+class FoldersProxyModel(RecursiveSortFilterProxyModel):
+    def __init__(self):
+        super().__init__()
+
+        self._folder_ids_filter = None
+
+    def set_folder_ids_filter(self, folder_ids: Optional[list[str]]):
+        if self._folder_ids_filter == folder_ids:
+            return
+        self._folder_ids_filter = folder_ids
+        self.invalidateFilter()
+
+    def filterAcceptsRow(self, row, parent_index):
+        if self._folder_ids_filter is not None:
+            if not self._folder_ids_filter:
+                return False
+            source_index = self.sourceModel().index(row, 0, parent_index)
+            folder_id = source_index.data(FOLDER_ID_ROLE)
+            if folder_id not in self._folder_ids_filter:
+                return False
+        return super().filterAcceptsRow(row, parent_index)
+
+
 class FoldersWidget(QtWidgets.QWidget):
     """Folders widget.
 
@@ -369,13 +393,13 @@ class FoldersWidget(QtWidgets.QWidget):
     refreshed = QtCore.Signal()
 
     def __init__(self, controller, parent, handle_expected_selection=False):
-        super(FoldersWidget, self).__init__(parent)
+        super().__init__(parent)
 
         folders_view = TreeView(self)
         folders_view.setHeaderHidden(True)
 
         folders_model = FoldersQtModel(controller)
-        folders_proxy_model = RecursiveSortFilterProxyModel()
+        folders_proxy_model = FoldersProxyModel()
         folders_proxy_model.setSourceModel(folders_model)
         folders_proxy_model.setSortCaseSensitivity(QtCore.Qt.CaseInsensitive)
 
@@ -446,6 +470,15 @@ class FoldersWidget(QtWidgets.QWidget):
         if name:
             self._folders_view.expandAll()
 
+    def set_folder_ids_filter(self, folder_ids: Optional[list[str]]):
+        """Set filter of folder ids.
+
+        Args:
+            folder_ids (list[str]): The list of folder ids.
+
+        """
+        self._folders_proxy_model.set_folder_ids_filter(folder_ids)
+
     def refresh(self):
         """Refresh folders model.
 
diff --git a/client/ayon_core/tools/utils/tasks_widget.py b/client/ayon_core/tools/utils/tasks_widget.py
index 30846e6cda..00960b6ba6 100644
--- a/client/ayon_core/tools/utils/tasks_widget.py
+++ b/client/ayon_core/tools/utils/tasks_widget.py
@@ -1,3 +1,5 @@
+from typing import Optional
+
 from qtpy import QtWidgets, QtGui, QtCore
 
 from ayon_core.style import (
@@ -343,6 +345,29 @@ class TasksQtModel(QtGui.QStandardItemModel):
         return self._has_content
 
 
+class TasksProxyModel(QtCore.QSortFilterProxyModel):
+    def __init__(self):
+        super().__init__()
+
+        self._task_ids_filter: Optional[set[str]] = None
+
+    def set_task_ids_filter(self, task_ids: Optional[set[str]]):
+        if self._task_ids_filter == task_ids:
+            return
+        self._task_ids_filter = task_ids
+        self.invalidateFilter()
+
+    def filterAcceptsRow(self, row, parent_index):
+        if self._task_ids_filter is not None:
+            if not self._task_ids_filter:
+                return False
+            source_index = self.sourceModel().index(row, 0, parent_index)
+            task_id = source_index.data(ITEM_ID_ROLE)
+            if task_id not in self._task_ids_filter:
+                return False
+        return super().filterAcceptsRow(row, parent_index)
+
+
 class TasksWidget(QtWidgets.QWidget):
     """Tasks widget.
 
@@ -364,7 +389,7 @@ class TasksWidget(QtWidgets.QWidget):
         tasks_view.setIndentation(0)
 
         tasks_model = TasksQtModel(controller)
-        tasks_proxy_model = QtCore.QSortFilterProxyModel()
+        tasks_proxy_model = TasksProxyModel()
         tasks_proxy_model.setSourceModel(tasks_model)
         tasks_proxy_model.setSortCaseSensitivity(QtCore.Qt.CaseInsensitive)
 
@@ -490,6 +515,15 @@ class TasksWidget(QtWidgets.QWidget):
         )
         return True
 
+    def set_task_ids_filter(self, task_ids: Optional[list[str]]):
+        """Set filter of folder ids.
+
+        Args:
+            task_ids (list[str]): The list of folder ids.
+
+        """
+        self._tasks_proxy_model.set_task_ids_filter(task_ids)
+
     def _on_tasks_refresh_finished(self, event):
         """Tasks were refreshed in controller.