use enum with named fields instead of ad-hoc numbers

Author: yedpodtrzitko

tagstudio/src/core/enums.py

@@ -14,3 +14,16 @@ class Theme(str, enum.Enum):
     COLOR_BG = "#65000000"
     COLOR_HOVER = "#65AAAAAA"
     COLOR_PRESSED = "#65EEEEEE"
+
+
+class FieldID(int, enum.Enum):
+    TITLE = 0
+    AUTHOR = 1
+    ARTIST = 2
+    DESCRIPTION = 4
+    NOTES = 5
+    TAGS = 6
+    CONTENT_TAGS = 7
+    META_TAGS = 8
+    DATE_PUBLISHED = 14
+    SOURCE = 21

tagstudio/src/core/library.py

@@ -16,6 +16,7 @@
 from enum import Enum
 import ujson
 
+from src.core.enums import FieldID
 from src.core.json_typing import JsonCollation, JsonEntry, JsonLibary, JsonTag
 from src.core import ts_core
 from src.core.utils.str import strip_punctuation
@@ -1960,48 +1961,44 @@ def add_generic_data_to_entry(self, data, entry_id: int):
         if data:
             # Add a Title Field if the data doesn't already exist.
             if data.get("title"):
-                field_id = 0  # Title Field ID
-                if not self.does_field_content_exist(entry_id, field_id, data["title"]):
-                    self.add_field_to_entry(entry_id, field_id)
+                if not self.does_field_content_exist(
+                    entry_id, FieldID.TITLE, data["title"]
+                ):
+                    self.add_field_to_entry(entry_id, FieldID.TITLE)
                     self.update_entry_field(entry_id, -1, data["title"], "replace")
 
             # Add an Author Field if the data doesn't already exist.
             if data.get("author"):
-                field_id = 1  # Author Field ID
                 if not self.does_field_content_exist(
-                    entry_id, field_id, data["author"]
+                    entry_id, FieldID.AUTHOR, data["author"]
                 ):
-                    self.add_field_to_entry(entry_id, field_id)
+                    self.add_field_to_entry(entry_id, FieldID.AUTHOR)
                     self.update_entry_field(entry_id, -1, data["author"], "replace")
 
             # Add an Artist Field if the data doesn't already exist.
             if data.get("artist"):
-                field_id = 2  # Artist Field ID
                 if not self.does_field_content_exist(
-                    entry_id, field_id, data["artist"]
+                    entry_id, FieldID.ARTIST, data["artist"]
                 ):
-                    self.add_field_to_entry(entry_id, field_id)
+                    self.add_field_to_entry(entry_id, FieldID.ARTIST)
                     self.update_entry_field(entry_id, -1, data["artist"], "replace")
 
             # Add a Date Published Field if the data doesn't already exist.
             if data.get("date_published"):
-                field_id = 14  # Date Published Field ID
                 date = str(
                     datetime.datetime.strptime(
                         data["date_published"], "%Y-%m-%d %H:%M:%S"
                     )
                 )
-                if not self.does_field_content_exist(entry_id, field_id, date):
-                    self.add_field_to_entry(entry_id, field_id)
+                if not self.does_field_content_exist(
+                    entry_id, FieldID.DATE_PUBLISHED, date
+                ):
+                    self.add_field_to_entry(entry_id, FieldID.DATE_PUBLISHED)
                     # entry = self.entries[entry_id]
                     self.update_entry_field(entry_id, -1, date, "replace")
 
             # Process String Tags if the data doesn't already exist.
             if data.get("tags"):
-                tags_field_id = 6  # Tags Field ID
-                content_tags_field_id = 7  # Content Tags Field ID
-                meta_tags_field_id = 8  # Meta Tags Field ID
-                notes_field_id = 5  # Notes Field ID
                 tags: list[str] = data["tags"]
                 # extra: list[str] = []
                 # for tag in tags:
@@ -2050,7 +2047,7 @@ def add_generic_data_to_entry(self, data, entry_id: int):
                         # tag_field_indices = self.get_field_index_in_entry(
                         # 	entry_index, tags_field_id)
                         content_tags_field_indices = self.get_field_index_in_entry(
-                            self.get_entry(entry_id), content_tags_field_id
+                            self.get_entry(entry_id), FieldID.CONTENT_TAGS
                         )
                         # meta_tags_field_indices = self.get_field_index_in_entry(
                         # 	entry_index, meta_tags_field_id)
@@ -2067,45 +2064,40 @@ def add_generic_data_to_entry(self, data, entry_id: int):
                                 entry_id, priority_field_index, [matching[0]], "append"
                             )
                         else:
-                            self.add_field_to_entry(entry_id, content_tags_field_id)
+                            self.add_field_to_entry(entry_id, FieldID.CONTENT_TAGS)
                             self.update_entry_field(
                                 entry_id, -1, [matching[0]], "append"
                             )
 
                 # Add all original string tags as a note.
                 str_tags = f"Original Tags: {tags}"
-                if not self.does_field_content_exist(
-                    entry_id, notes_field_id, str_tags
-                ):
-                    self.add_field_to_entry(entry_id, notes_field_id)
+                if not self.does_field_content_exist(entry_id, FieldID.NOTES, str_tags):
+                    self.add_field_to_entry(entry_id, FieldID.NOTES)
                     self.update_entry_field(entry_id, -1, str_tags, "replace")
 
             # Add a Description Field if the data doesn't already exist.
-            if "description" in data.keys() and data["description"]:
-                field_id = 4  # Description Field ID
+            if data.get("description"):
                 if not self.does_field_content_exist(
-                    entry_id, field_id, data["description"]
+                    entry_id, FieldID.DESCRIPTION, data["description"]
                 ):
-                    self.add_field_to_entry(entry_id, field_id)
+                    self.add_field_to_entry(entry_id, FieldID.DESCRIPTION)
                     self.update_entry_field(
                         entry_id, -1, data["description"], "replace"
                     )
-            if "content" in data.keys() and data["content"]:
-                field_id = 4  # Description Field ID
+            if data.get("content"):
                 if not self.does_field_content_exist(
-                    entry_id, field_id, data["content"]
+                    entry_id, FieldID.DESCRIPTION, data["content"]
                 ):
-                    self.add_field_to_entry(entry_id, field_id)
+                    self.add_field_to_entry(entry_id, FieldID.DESCRIPTION)
                     self.update_entry_field(entry_id, -1, data["content"], "replace")
-            if "source" in data.keys() and data["source"]:
-                field_id = 21  # Source Field ID
+            if data.get("source"):
                 for source in data["source"].split(" "):
                     if source and source != " ":
                         source = strip_web_protocol(string=source)
                         if not self.does_field_content_exist(
-                            entry_id, field_id, source
+                            entry_id, FieldID.SOURCE, source
                         ):
-                            self.add_field_to_entry(entry_id, field_id)
+                            self.add_field_to_entry(entry_id, FieldID.SOURCE)
                             self.update_entry_field(entry_id, -1, source, "replace")
 
     def add_field_to_entry(self, entry_id: int, field_id: int) -> None:

tagstudio/src/core/ts_core.py

@@ -6,6 +6,7 @@
 
 import json
 import os
+from enum import Enum
 
 from src.core.library import Entry, Library
 
@@ -18,6 +19,7 @@
 COLLAGE_FOLDER_NAME: str = "collages"
 LIBRARY_FILENAME: str = "ts_library.json"
 
+
 # TODO: Turn this whitelist into a user-configurable blacklist.
 IMAGE_TYPES: list[str] = [
     "png",

tagstudio/src/qt/modals/folders_to_tags.py

@@ -18,6 +18,7 @@
     QFrame,
 )
 
+from src.core.enums import FieldID
 from src.core.library import Library, Tag
 from src.core.palette import ColorType, get_tag_color
 from src.qt.flowlayout import FlowLayout
@@ -73,7 +74,7 @@ def add_folders_to_tree(list: list[str]) -> Tag:
         tag = add_folders_to_tree(folders)
         if tag:
             if not entry.has_tag(library, tag.id):
-                entry.add_tag(library, tag.id, 6)
+                entry.add_tag(library, tag.id, FieldID.TAGS)
 
     logging.info("Done")
 

tagstudio/src/qt/ts_qt.py

@@ -1195,8 +1195,8 @@ def update_thumbs(self):
                     )
 
                     item_thumb.set_item_id(entry.id)
-                    item_thumb.assign_archived(entry.has_tag(self.lib, 0))
-                    item_thumb.assign_favorite(entry.has_tag(self.lib, 1))
+                    item_thumb.assign_archived(entry.has_tag(self.lib, TAG_ARCHIVED))
+                    item_thumb.assign_favorite(entry.has_tag(self.lib, TAG_FAVORITE))
                     # ctrl_down = True if QGuiApplication.keyboardModifiers() else False
                     # TODO: Change how this works. The click function
                     # for collations a few lines down should NOT be allowed during modifier keys.

tagstudio/src/qt/widgets/item_thumb.py

@@ -23,6 +23,7 @@
     QCheckBox,
 )
 
+from src.core.enums import FieldID
 from src.core.library import ItemType, Library, Entry
 from src.core.ts_core import AUDIO_TYPES, VIDEO_TYPES, IMAGE_TYPES
 from src.qt.flowlayout import FlowWidget
@@ -37,7 +38,6 @@
 WARNING = f"[WARNING]"
 INFO = f"[INFO]"
 
-DEFAULT_META_TAG_FIELD = 8
 TAG_FAVORITE = 1
 TAG_ARCHIVED = 0
 
@@ -402,8 +402,12 @@ def update_badges(self):
         if self.mode == ItemType.ENTRY:
             # logging.info(f'[UPDATE BADGES] ENTRY: {self.lib.get_entry(self.item_id)}')
             # logging.info(f'[UPDATE BADGES] ARCH: {self.lib.get_entry(self.item_id).has_tag(self.lib, 0)}, FAV: {self.lib.get_entry(self.item_id).has_tag(self.lib, 1)}')
-            self.assign_archived(self.lib.get_entry(self.item_id).has_tag(self.lib, 0))
-            self.assign_favorite(self.lib.get_entry(self.item_id).has_tag(self.lib, 1))
+            self.assign_archived(
+                self.lib.get_entry(self.item_id).has_tag(self.lib, TAG_ARCHIVED)
+            )
+            self.assign_favorite(
+                self.lib.get_entry(self.item_id).has_tag(self.lib, TAG_FAVORITE)
+            )
 
     def set_item_id(self, id: int):
         """
@@ -474,7 +478,7 @@ def toggle_tag(entry: Entry):
                 entry.add_tag(
                     self.panel.driver.lib,
                     tag_id,
-                    field_id=DEFAULT_META_TAG_FIELD,
+                    field_id=FieldID.META_TAGS,
                     field_index=-1,
                 )
             else:

tagstudio/src/qt/widgets/tag_box.py

@@ -142,7 +142,7 @@ def edit_tag(self, tag_id: int):
         # panel.tag_updated.connect(lambda tag: self.lib.update_tag(tag))
         self.edit_modal.show()
 
-    def add_tag_callback(self, tag_id):
+    def add_tag_callback(self, tag_id: int):
         # self.base_layout.addWidget(TagWidget(self.lib, self.lib.get_tag(tag), True))
         # self.tags.append(tag)
         logging.info(
@@ -155,7 +155,7 @@ def add_tag_callback(self, tag_id):
                 self.driver.lib, tag_id, field_id=id, field_index=-1
             )
             self.updated.emit()
-        if tag_id == 0 or tag_id == 1:
+        if tag_id in (TAG_FAVORITE, TAG_ARCHIVED):
             self.driver.update_badges()
 
         # if type((x[0]) == ThumbButton):
@@ -170,7 +170,7 @@ def add_tag_callback(self, tag_id):
     def edit_tag_callback(self, tag: Tag):
         self.lib.update_tag(tag)
 
-    def remove_tag(self, tag_id):
+    def remove_tag(self, tag_id: int):
         logging.info(f"[TAG BOX WIDGET] SELECTED T:{self.driver.selected}")
         id = list(self.field.keys())[0]
         for x in self.driver.selected:
@@ -181,7 +181,7 @@ def remove_tag(self, tag_id):
                 self.driver.lib, tag_id, field_index=index[0]
             )
             self.updated.emit()
-        if tag_id == 0 or tag_id == 1:
+        if tag_id in (TAG_FAVORITE, TAG_ARCHIVED):
             self.driver.update_badges()
 
     # def show_add_button(self, value:bool):