From b6b24b698c25534cd0216fea2bb941c9bc3192f6 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Tue, 20 Jan 2026 11:57:10 +0100 Subject: [PATCH] Update engine.py --- engine.py | 153 ++++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 150 insertions(+), 3 deletions(-) diff --git a/engine.py b/engine.py index e6ee1ca..5bc4bbc 100644 --- a/engine.py +++ b/engine.py @@ -48,7 +48,7 @@ class SorterEngine: @staticmethod def init_db(): - """Initializes tables, including the HISTORY log.""" + """Initializes tables, including the HISTORY log and persistent tags.""" with SorterEngine.get_connection() as conn: cursor = conn.cursor() @@ -62,6 +62,21 @@ class SorterEngine: cursor.execute('''CREATE TABLE IF NOT EXISTS processed_log (source_path TEXT PRIMARY KEY, category TEXT, action_type TEXT)''') + # NEW: Persistent tags table - survives after commit + # Maps: output_folder + filename -> original source info for tag restoration + cursor.execute('''CREATE TABLE IF NOT EXISTS persistent_tags + (output_folder TEXT, + filename TEXT, + category TEXT, + tag_index INTEGER, + original_source_folder TEXT, + committed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (output_folder, filename))''') + + # Index for fast lookup by output folder + cursor.execute('''CREATE INDEX IF NOT EXISTS idx_persistent_tags_folder + ON persistent_tags(output_folder)''') + # Seed categories if empty cursor.execute("SELECT COUNT(*) FROM categories") if cursor.fetchone()[0] == 0: @@ -388,6 +403,88 @@ class SorterEngine: rows = cursor.fetchall() return {r[0]: {"cat": r[1], "name": r[2], "marked": r[3]} for r in rows} + # ========================================== + # PERSISTENT TAGS (NEW) + # ========================================== + + @staticmethod + def save_persistent_tag(output_folder: str, filename: str, category: str, + tag_index: int, source_folder: str): + """Save a tag permanently after commit.""" + with SorterEngine.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + INSERT OR REPLACE INTO persistent_tags + (output_folder, filename, category, tag_index, original_source_folder) + VALUES (?, ?, ?, ?, ?) + """, (output_folder, filename, category, tag_index, source_folder)) + conn.commit() + + @staticmethod + def save_persistent_tags_batch(tags: List[Tuple[str, str, str, int, str]]): + """Batch save multiple persistent tags.""" + with SorterEngine.transaction() as conn: + cursor = conn.cursor() + cursor.executemany(""" + INSERT OR REPLACE INTO persistent_tags + (output_folder, filename, category, tag_index, original_source_folder) + VALUES (?, ?, ?, ?, ?) + """, tags) + + @staticmethod + def get_persistent_tags(output_folder: str) -> Dict[str, Dict]: + """ + Get all persistent tags for an output folder. + Returns: {filename: {"cat": category, "index": tag_index, "source": source_folder}} + """ + with SorterEngine.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + SELECT filename, category, tag_index, original_source_folder + FROM persistent_tags + WHERE output_folder = ? + """, (output_folder,)) + rows = cursor.fetchall() + return { + r[0]: {"cat": r[1], "index": r[2], "source": r[3]} + for r in rows + } + + @staticmethod + def get_persistent_tags_by_category(output_folder: str, category: str) -> Dict[int, str]: + """ + Get persistent tags for a specific category in an output folder. + Returns: {tag_index: filename} + """ + with SorterEngine.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + SELECT tag_index, filename + FROM persistent_tags + WHERE output_folder = ? AND category = ? + """, (output_folder, category)) + rows = cursor.fetchall() + return {r[0]: r[1] for r in rows} + + @staticmethod + def delete_persistent_tag(output_folder: str, filename: str): + """Remove a persistent tag (e.g., if file is deleted).""" + with SorterEngine.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + DELETE FROM persistent_tags + WHERE output_folder = ? AND filename = ? + """, (output_folder, filename)) + conn.commit() + + @staticmethod + def clear_persistent_tags(output_folder: str): + """Clear all persistent tags for an output folder.""" + with SorterEngine.get_connection() as conn: + cursor = conn.cursor() + cursor.execute("DELETE FROM persistent_tags WHERE output_folder = ?", (output_folder,)) + conn.commit() + # ========================================== # BATCH COMMIT OPERATIONS (OPTIMIZED) # ========================================== @@ -416,7 +513,7 @@ class SorterEngine: @staticmethod def commit_batch(file_list: List[str], output_root: str, cleanup_mode: str, operation: str = "Copy"): - """Commits files with batched DB operations.""" + """Commits files with batched DB operations and saves persistent tags.""" data = SorterEngine.get_staged_data() if not os.path.exists(output_root): @@ -425,6 +522,7 @@ class SorterEngine: # Prepare batch operations to_delete_from_staging = [] to_insert_to_log = [] + persistent_tags_to_save = [] for file_path in file_list: if not os.path.exists(file_path): @@ -434,6 +532,14 @@ class SorterEngine: if file_path in data and data[file_path]['marked']: info = data[file_path] final_dst = SorterEngine._compute_final_destination(output_root, info['name']) + final_filename = os.path.basename(final_dst) + + # Extract tag index from filename + tag_index = None + try: + tag_index = int(info['name'].rsplit('_', 1)[1].split('.')[0]) + except (ValueError, IndexError): + pass # Perform file operation if operation == "Copy": @@ -445,6 +551,16 @@ class SorterEngine: to_delete_from_staging.append((file_path,)) to_insert_to_log.append((file_path, info['cat'], operation)) + + # Save persistent tag + if tag_index is not None: + persistent_tags_to_save.append(( + output_root, # output_folder + final_filename, # filename + info['cat'], # category + tag_index, # tag_index + os.path.dirname(file_path) # original source folder + )) # Untagged files - cleanup elif cleanup_mode != "Keep": @@ -464,16 +580,23 @@ class SorterEngine: cursor.executemany("DELETE FROM staging_area WHERE original_path = ?", to_delete_from_staging) if to_insert_to_log: cursor.executemany("INSERT OR REPLACE INTO processed_log VALUES (?, ?, ?)", to_insert_to_log) + if persistent_tags_to_save: + cursor.executemany(""" + INSERT OR REPLACE INTO persistent_tags + (output_folder, filename, category, tag_index, original_source_folder) + VALUES (?, ?, ?, ?, ?) + """, persistent_tags_to_save) @staticmethod def commit_global(output_root: str, cleanup_mode: str, operation: str = "Copy", source_root: str = None): - """Commits ALL staged files with batched operations.""" + """Commits ALL staged files with batched operations and saves persistent tags.""" data = SorterEngine.get_staged_data() if not os.path.exists(output_root): os.makedirs(output_root, exist_ok=True) to_insert_to_log = [] + persistent_tags_to_save = [] # Process all staged items for old_p, info in data.items(): @@ -481,6 +604,14 @@ class SorterEngine: continue final_dst = SorterEngine._compute_final_destination(output_root, info['name']) + final_filename = os.path.basename(final_dst) + + # Extract tag index from filename + tag_index = None + try: + tag_index = int(info['name'].rsplit('_', 1)[1].split('.')[0]) + except (ValueError, IndexError): + pass if operation == "Copy": shutil.copy2(old_p, final_dst) @@ -489,6 +620,16 @@ class SorterEngine: SorterEngine.fix_permissions(final_dst) to_insert_to_log.append((old_p, info['cat'], operation)) + + # Save persistent tag + if tag_index is not None: + persistent_tags_to_save.append(( + output_root, # output_folder + final_filename, # filename + info['cat'], # category + tag_index, # tag_index + os.path.dirname(old_p) # original source folder + )) # Global cleanup if cleanup_mode != "Keep" and source_root: @@ -510,6 +651,12 @@ class SorterEngine: cursor.execute("DELETE FROM staging_area") if to_insert_to_log: cursor.executemany("INSERT OR REPLACE INTO processed_log VALUES (?, ?, ?)", to_insert_to_log) + if persistent_tags_to_save: + cursor.executemany(""" + INSERT OR REPLACE INTO persistent_tags + (output_folder, filename, category, tag_index, original_source_folder) + VALUES (?, ?, ?, ?, ?) + """, persistent_tags_to_save) # ========================================== # UTILITY OPERATIONS