Update engine.py
This commit is contained in:
120
engine.py
120
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
|
||||
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
|
||||
Reference in New Issue
Block a user