diff --git a/engine.py b/engine.py index c43446b..55c0439 100644 --- a/engine.py +++ b/engine.py @@ -27,6 +27,11 @@ class SorterEngine: cursor.execute('''CREATE TABLE IF NOT EXISTS processed_log (source_path TEXT PRIMARY KEY, category TEXT, action_type TEXT)''') + # --- NEW: FOLDER TAGS TABLE (persists tags by folder) --- + cursor.execute('''CREATE TABLE IF NOT EXISTS folder_tags + (folder_path TEXT, filename TEXT, category TEXT, tag_index INTEGER, + PRIMARY KEY (folder_path, filename))''') + # Seed categories if empty cursor.execute("SELECT COUNT(*) FROM categories") if cursor.fetchone()[0] == 0: @@ -256,6 +261,11 @@ class SorterEngine: def commit_global(output_root, cleanup_mode, operation="Copy", source_root=None): """Commits ALL staged files and fixes permissions.""" data = SorterEngine.get_staged_data() + + # Save folder tags BEFORE processing (so we can restore them later) + if source_root: + SorterEngine.save_folder_tags(source_root) + conn = sqlite3.connect(SorterEngine.DB_PATH) cursor = conn.cursor() @@ -513,4 +523,112 @@ class SorterEngine: for idx, img_path in enumerate(all_images): if img_path in staged_keys: tagged_pages.add(idx // page_size) - return tagged_pages \ No newline at end of file + return tagged_pages + + # --- 7. FOLDER TAG PERSISTENCE --- + @staticmethod + def save_folder_tags(folder_path): + """ + Saves current staging data associated with a folder for later restoration. + Call this BEFORE clearing the staging area. + """ + import re + staged = SorterEngine.get_staged_data() + if not staged: + return 0 + + conn = sqlite3.connect(SorterEngine.DB_PATH) + cursor = conn.cursor() + + saved_count = 0 + for orig_path, info in staged.items(): + # Only save tags for files that are in this folder (or subfolders) + if orig_path.startswith(folder_path): + filename = os.path.basename(orig_path) + category = info['cat'] + + # Extract index from the new_name (e.g., "Action_042.jpg" -> 42) + new_name = info['name'] + match = re.search(r'_(\d+)', new_name) + tag_index = int(match.group(1)) if match else 0 + + cursor.execute( + "INSERT OR REPLACE INTO folder_tags VALUES (?, ?, ?, ?)", + (folder_path, filename, category, tag_index) + ) + saved_count += 1 + + conn.commit() + conn.close() + return saved_count + + @staticmethod + def restore_folder_tags(folder_path, all_images): + """ + Restores previously saved tags for a folder back into the staging area. + Call this when loading/reloading a folder. + Returns the number of tags restored. + """ + conn = sqlite3.connect(SorterEngine.DB_PATH) + cursor = conn.cursor() + + # Get saved tags for this folder + cursor.execute( + "SELECT filename, category, tag_index FROM folder_tags WHERE folder_path = ?", + (folder_path,) + ) + saved_tags = {row[0]: {"cat": row[1], "index": row[2]} for row in cursor.fetchall()} + + if not saved_tags: + conn.close() + return 0 + + # Build a map of filename -> full path from current images + filename_to_path = {} + for img_path in all_images: + fname = os.path.basename(img_path) + # Only map files that haven't been mapped yet (handles duplicates by using first occurrence) + if fname not in filename_to_path: + filename_to_path[fname] = img_path + + # Restore tags to staging area + restored = 0 + for filename, tag_info in saved_tags.items(): + if filename in filename_to_path: + full_path = filename_to_path[filename] + # Check if not already staged + cursor.execute("SELECT 1 FROM staging_area WHERE original_path = ?", (full_path,)) + if not cursor.fetchone(): + ext = os.path.splitext(filename)[1] + new_name = f"{tag_info['cat']}_{tag_info['index']:03d}{ext}" + cursor.execute( + "INSERT OR REPLACE INTO staging_area VALUES (?, ?, ?, 1)", + (full_path, tag_info['cat'], new_name) + ) + restored += 1 + + conn.commit() + conn.close() + return restored + + @staticmethod + def clear_folder_tags(folder_path): + """Clears saved tags for a specific folder.""" + conn = sqlite3.connect(SorterEngine.DB_PATH) + cursor = conn.cursor() + cursor.execute("DELETE FROM folder_tags WHERE folder_path = ?", (folder_path,)) + conn.commit() + conn.close() + + @staticmethod + def get_saved_folder_tags(folder_path): + """Returns saved tags for a folder (for debugging/display).""" + conn = sqlite3.connect(SorterEngine.DB_PATH) + cursor = conn.cursor() + cursor.execute( + "SELECT filename, category, tag_index FROM folder_tags WHERE folder_path = ?", + (folder_path,) + ) + result = {row[0]: {"cat": row[1], "index": row[2]} for row in cursor.fetchall()} + conn.close() + return result \ No newline at end of file