Multi mode search system

tagstudio/src/core/enums.py

@@ -17,3 +17,10 @@ class Theme(str, enum.Enum):
     COLOR_PRESSED = "#65EEEEEE"
     COLOR_DISABLED = "#65F39CAA"
     COLOR_DISABLED_BG = "#65440D12"
+
+
+class SearchMode(int, enum.Enum):
+    """Operational modes for item searching."""
+
+    AND = 0
+    OR = 1

tagstudio/src/core/library.py

@@ -20,6 +20,7 @@
 from src.core.json_typing import JsonCollation, JsonEntry, JsonLibary, JsonTag
 from src.core.utils.str import strip_punctuation
 from src.core.utils.web import strip_web_protocol
+from src.core.enums import SearchMode
 from src.core.constants import (
     BACKUP_FOLDER_NAME,
     COLLAGE_FOLDER_NAME,
@@ -1290,7 +1291,12 @@ def get_entry_id_from_filepath(self, filename: Path):
             return -1
 
     def search_library(
-        self, query: str = None, entries=True, collations=True, tag_groups=True
+        self,
+        query: str = None,
+        entries=True,
+        collations=True,
+        tag_groups=True,
+        search_mode=SearchMode.AND,
     ) -> list[tuple[ItemType, int]]:
         """
         Uses a search query to generate a filtered results list.
@@ -1300,7 +1306,7 @@ def search_library(
         # self.filtered_entries.clear()
         results: list[tuple[ItemType, int]] = []
         collations_added = []
-
+        # print(f"Searching Library with query: {query} search_mode: {search_mode}")
         if query:
             # start_time = time.time()
             query = query.strip().lower()
@@ -1320,6 +1326,7 @@ def search_library(
 
             # Preprocess the Tag terms.
             if query_words:
+                # print(query_words, self._tag_strings_to_id_map)
                 for i, term in enumerate(query_words):
                     for j, term in enumerate(query_words):
                         if (
@@ -1328,6 +1335,8 @@ def search_library(
                             in self._tag_strings_to_id_map
                         ):
                             all_tag_terms.append(" ".join(query_words[i : j + 1]))
+                        # print(all_tag_terms)
+
                 # This gets rid of any accidental term inclusions because they were words
                 # in another term. Ex. "3d" getting added in "3d art"
                 for i, term in enumerate(all_tag_terms):
@@ -1403,36 +1412,8 @@ def search_library(
                     # elif query in entry.filename.lower():
                     # 	self.filtered_entries.append(index)
                     elif entry_tags:
-                        # For each verified, extracted Tag term.
-                        failure_to_union_terms = False
-                        for term in all_tag_terms:
-                            # If the term from the previous loop was already verified:
-                            if not failure_to_union_terms:
-                                cluster: set = set()
-                                # Add the immediate associated Tags to the set (ex. Name, Alias hits)
-                                # Since this term could technically map to multiple IDs, iterate over it
-                                # (You're 99.9999999% likely to just get 1 item)
-                                for id in self._tag_strings_to_id_map[term]:
-                                    cluster.add(id)
-                                    cluster = cluster.union(
-                                        set(self.get_tag_cluster(id))
-                                    )
-                                # print(f'Full Cluster: {cluster}')
-                                # For each of the Tag IDs in the term's ID cluster:
-                                for t in cluster:
-                                    # Assume that this ID from the cluster is not in the Entry.
-                                    # Wait to see if proven wrong.
-                                    failure_to_union_terms = True
-                                    # If the ID actually is in the Entry,
-                                    if t in entry_tags:
-                                        # There wasn't a failure to find one of the term's cluster IDs in the Entry.
-                                        # There is also no more need to keep checking the rest of the terms in the cluster.
-                                        failure_to_union_terms = False
-                                        # print(f'FOUND MATCH: {t}')
-                                        break
-                                    # print(f'\tFailure to Match: {t}')
-                        # If there even were tag terms to search through AND they all match an entry
-                        if all_tag_terms and not failure_to_union_terms:
+                        # function to add entry to results
+                        def add_entry(entry: Entry):
                             # self.filter_entries.append()
                             # self.filtered_file_list.append(file)
                             # results.append((SearchItemType.ENTRY, entry.id))
@@ -1457,6 +1438,54 @@ def search_library(
                             if not added:
                                 results.append((ItemType.ENTRY, entry.id))
 
+                        if search_mode == SearchMode.AND:  # Include all terms
+                            # For each verified, extracted Tag term.
+                            failure_to_union_terms = False
+                            for term in all_tag_terms:
+                                # If the term from the previous loop was already verified:
+                                if not failure_to_union_terms:
+                                    cluster: set = set()
+                                    # Add the immediate associated Tags to the set (ex. Name, Alias hits)
+                                    # Since this term could technically map to multiple IDs, iterate over it
+                                    # (You're 99.9999999% likely to just get 1 item)
+                                    for id in self._tag_strings_to_id_map[term]:
+                                        cluster.add(id)
+                                        cluster = cluster.union(
+                                            set(self.get_tag_cluster(id))
+                                        )
+                                    # print(f'Full Cluster: {cluster}')
+                                    # For each of the Tag IDs in the term's ID cluster:
+                                    for t in cluster:
+                                        # Assume that this ID from the cluster is not in the Entry.
+                                        # Wait to see if proven wrong.
+                                        failure_to_union_terms = True
+                                        # If the ID actually is in the Entry,
+                                        if t in entry_tags:
+                                            # There wasn't a failure to find one of the term's cluster IDs in the Entry.
+                                            # There is also no more need to keep checking the rest of the terms in the cluster.
+                                            failure_to_union_terms = False
+                                            # print(f"FOUND MATCH: {t}")
+                                            break
+                                        # print(f'\tFailure to Match: {t}')
+                            # # failure_to_union_terms is used to determine if all terms in the query were found in the entry.
+                            # # If there even were tag terms to search through AND they all match an entry
+                            if all_tag_terms and not failure_to_union_terms:
+                                add_entry(entry)
+
+                        if search_mode == SearchMode.OR:  # Include any terms
+                            # For each verified, extracted Tag term.
+                            for term in all_tag_terms:
+                                # Add the immediate associated Tags to the set (ex. Name, Alias hits)
+                                # Since this term could technically map to multiple IDs, iterate over it
+                                # (You're 99.9999999% likely to just get 1 item)
+                                for id in self._tag_strings_to_id_map[term]:
+                                    # If the ID actually is in the Entry,
+                                    if id in entry_tags:
+                                        # check if result already contains the entry
+                                        if (ItemType.ENTRY, entry.id) not in results:
+                                            add_entry(entry)
+                                        break
+
                 # sys.stdout.write(
                 #     f'\r[INFO][FILTER]: {len(self.filtered_file_list)} matches found')
                 # sys.stdout.flush()

tagstudio/src/qt/main_window.py

@@ -18,7 +18,8 @@
 from PySide6.QtWidgets import (QComboBox, QFrame, QGridLayout,
 							   QHBoxLayout, QVBoxLayout, QLayout, QLineEdit, QMainWindow,
 							   QPushButton, QScrollArea, QSizePolicy,
-							   QStatusBar, QWidget, QSplitter)
+							   QStatusBar, QWidget, QSplitter, QCheckBox,
+							   QSpacerItem)
 from src.qt.pagination import Pagination
 
 
@@ -52,6 +53,36 @@ def setupUi(self, MainWindow):
 		self.gridLayout.setObjectName(u"gridLayout")
 		self.horizontalLayout = QHBoxLayout()
 		self.horizontalLayout.setObjectName(u"horizontalLayout")
+
+		# ComboBox goup for search type and thumbnail size
+		self.horizontalLayout_3 = QHBoxLayout()
+		self.horizontalLayout_3.setObjectName("horizontalLayout_3")
+
+		# left side spacer
+		spacerItem = QSpacerItem(40, 20, QSizePolicy.Expanding, QSizePolicy.Minimum)
+		self.horizontalLayout_3.addItem(spacerItem)
+
+  		# Search type selector
+		self.comboBox_2 = QComboBox(self.centralwidget)
+		self.comboBox_2.setMinimumSize(QSize(165, 0))
+		self.comboBox_2.setObjectName("comboBox_2")
+		self.comboBox_2.addItem("")
+		self.comboBox_2.addItem("")
+		self.horizontalLayout_3.addWidget(self.comboBox_2)
+
+        # Thumbnail Size placeholder
+		self.comboBox = QComboBox(self.centralwidget)
+		self.comboBox.setObjectName(u"comboBox")
+		sizePolicy = QSizePolicy(QSizePolicy.Minimum, QSizePolicy.Fixed)
+		sizePolicy.setHorizontalStretch(0)
+		sizePolicy.setVerticalStretch(0)
+		sizePolicy.setHeightForWidth(
+			self.comboBox.sizePolicy().hasHeightForWidth())
+		self.comboBox.setSizePolicy(sizePolicy)
+		self.comboBox.setMinimumWidth(128)
+		self.comboBox.setMaximumWidth(128)
+		self.horizontalLayout_3.addWidget(self.comboBox)
+		self.gridLayout.addLayout(self.horizontalLayout_3, 5, 0, 1, 1)
 
 		self.splitter = QSplitter()
 		self.splitter.setObjectName(u"splitter")
@@ -138,18 +169,18 @@ def setupUi(self, MainWindow):
 		self.horizontalLayout_2.addWidget(self.searchButton)
 		self.gridLayout.addLayout(self.horizontalLayout_2, 3, 0, 1, 1)
 
-		self.comboBox = QComboBox(self.centralwidget)
-		self.comboBox.setObjectName(u"comboBox")
-		sizePolicy = QSizePolicy(QSizePolicy.Minimum, QSizePolicy.Fixed)
-		sizePolicy.setHorizontalStretch(0)
-		sizePolicy.setVerticalStretch(0)
-		sizePolicy.setHeightForWidth(
-			self.comboBox.sizePolicy().hasHeightForWidth())
-		self.comboBox.setSizePolicy(sizePolicy)
-		self.comboBox.setMinimumWidth(128)
-		self.comboBox.setMaximumWidth(128)
+		# self.comboBox = QComboBox(self.centralwidget)
+		# self.comboBox.setObjectName(u"comboBox")
+		# sizePolicy = QSizePolicy(QSizePolicy.Minimum, QSizePolicy.Fixed)
+		# sizePolicy.setHorizontalStretch(0)
+		# sizePolicy.setVerticalStretch(0)
+		# sizePolicy.setHeightForWidth(
+		# 	self.comboBox.sizePolicy().hasHeightForWidth())
+		# self.comboBox.setSizePolicy(sizePolicy)
+		# self.comboBox.setMinimumWidth(128)
+		# self.comboBox.setMaximumWidth(128)
 
-		self.gridLayout.addWidget(self.comboBox, 4, 0, 1, 1, Qt.AlignRight)
+		# self.gridLayout.addWidget(self.comboBox, 4, 0, 1, 1, Qt.AlignRight)
 
 		self.gridLayout_2.setContentsMargins(6, 6, 6, 6)
 
@@ -181,15 +212,24 @@ def setupUi(self, MainWindow):
 	def retranslateUi(self, MainWindow):
 		MainWindow.setWindowTitle(QCoreApplication.translate(
 			"MainWindow", u"MainWindow", None))
+        # Navigation buttons
 		self.backButton.setText(
 			QCoreApplication.translate("MainWindow", u"<", None))
 		self.forwardButton.setText(
 			QCoreApplication.translate("MainWindow", u">", None))
+
+        # Search field
 		self.searchField.setPlaceholderText(
 			QCoreApplication.translate("MainWindow", u"Search Entries", None))
 		self.searchButton.setText(
 			QCoreApplication.translate("MainWindow", u"Search", None))
+
+        # Search type selector
+		self.comboBox_2.setItemText(0, QCoreApplication.translate("MainWindow", "And (Includes All Tags)"))
+		self.comboBox_2.setItemText(1, QCoreApplication.translate("MainWindow", "Or (Includes Any Tag)"))
 		self.comboBox.setCurrentText("")
+
+        # Tumbnail size selector
 		self.comboBox.setPlaceholderText(
 			QCoreApplication.translate("MainWindow", u"Thumbnail Size", None))
 	# retranslateUi

tagstudio/src/qt/ts_qt.py

@@ -43,10 +43,11 @@
     QSplashScreen,
     QMenu,
     QMenuBar,
+    QComboBox,
 )
 from humanfriendly import format_timespan
 
-from src.core.enums import SettingItems
+from src.core.enums import SettingItems, SearchMode
 from src.core.library import ItemType
 from src.core.ts_core import TagStudioCore
 from src.core.constants import (
@@ -176,6 +177,8 @@ def __init__(self, core: TagStudioCore, args):
         self.nav_frames: list[NavigationState] = []
         self.cur_frame_idx: int = -1
 
+        self.search_mode = SearchMode.AND
+
         # self.main_window = None
         # self.main_window = Ui_MainWindow()
 
@@ -564,6 +567,12 @@ def init_library_window(self):
         search_field.returnPressed.connect(
             lambda: self.filter_items(self.main_window.searchField.text())
         )
+        search_type_selector: QComboBox = self.main_window.comboBox_2
+        search_type_selector.currentIndexChanged.connect(
+            lambda: self.set_search_type(
+                SearchMode(search_type_selector.currentIndex())
+            )
+        )
 
         back_button: QPushButton = self.main_window.backButton
         back_button.clicked.connect(self.nav_back)
@@ -1334,7 +1343,7 @@ def filter_items(self, query: str = ""):
 
             # self.filtered_items = self.lib.search_library(query)
             # 73601 Entries at 500 size should be 246
-            all_items = self.lib.search_library(query)
+            all_items = self.lib.search_library(query, search_mode=self.search_mode)
             frames: list[list[tuple[ItemType, int]]] = []
             frame_count = math.ceil(len(all_items) / self.max_results)
             for i in range(0, frame_count):
@@ -1375,6 +1384,10 @@ def filter_items(self, query: str = ""):
 
             # self.update_thumbs()
 
+    def set_search_type(self, mode=SearchMode.AND):
+        self.search_mode = mode
+        self.filter_items(self.main_window.searchField.text())
+
     def remove_recent_library(self, item_key: str):
         self.settings.beginGroup(SettingItems.LIBS_LIST)
         self.settings.remove(item_key)