Cancellable tasks

kapitanbooru_uploader/ImageBrowser.py

@@ -6,7 +6,7 @@ import queue
 import threading
 import tkinter as tk
 from tkinter import filedialog, messagebox, ttk
-from typing import Tuple
+from typing import Dict, Tuple, Optional
 
 import networkx as nx
 import requests
@@ -41,6 +41,7 @@ class ProcessingDialog:
         # Setup communication queue and periodic checker
         self.queue = queue.Queue()
         self.running = True
+        self.cancel_event = threading.Event()  # Cancellation flag
         self.thread = threading.Thread(
             target=self.run_task, args=(target_function, *args)
         )
@@ -75,10 +76,12 @@ class ProcessingDialog:
         """Execute target function with progress queue if supported"""
         try:
             sig = inspect.signature(target_function)
+            kwargs = {}
             if "progress_queue" in sig.parameters:
-                target_function(*args, progress_queue=self.queue)
-            else:
-                target_function(*args)
+                kwargs["progress_queue"] = self.queue
+            if "cancel_event" in sig.parameters:
+                kwargs["cancel_event"] = self.cancel_event
+            target_function(*args, **kwargs)
         finally:
             self.close_dialog()
 
@@ -88,10 +91,12 @@ class ProcessingDialog:
             self.running = False
             self.progress.stop()
             self.top.after(0, self.top.destroy)
+        self.thread.join()
 
     def on_close(self):
         """Handle manual window closure"""
         self.running = False
+        self.cancel_event.set()  # Notify target function that cancellation is requested
         self.top.destroy()
 
 
@@ -121,6 +126,14 @@ class ImageBrowser(tk.Tk):
         self.image_files_md5 = []
         self.current_index = None
         self.image_cache = None
+        self.tagger_thread_idx = 0
+        self.tagger = wdt.Tagger()
+
+        self.check_uploaded_files_stop_event = threading.Event()
+        self.check_uploaded_files_thread: Optional[threading.Thread] = None
+        self.process_tagger_queue_stop_event = threading.Event()
+        self.process_tagger_queue_thread: Optional[threading.Thread] = None
+        self.run_tagger_threads: Dict[str, threading.Thread] = {}
 
         # Liczniki statusu
         self.total_files = 0
@@ -273,9 +286,8 @@ class ImageBrowser(tk.Tk):
             self.tagger_processed.add(md5)
             return cached["result"]
         try:
-            tagger = wdt.Tagger()
             with Image.open(file_path) as img:
-                result = tagger.tag(img)
+                result = self.tagger.tag(img)
             self.tagger_cache[md5] = result
             self.tagger_processed.add(md5)
             self.after(0, self.update_status_bar)
@@ -407,18 +419,6 @@ class ImageBrowser(tk.Tk):
         btn_save.pack(pady=10)
 
     def create_widgets(self):
-        # Górna ramka
-        top_frame = tk.Frame(self)
-        top_frame.pack(side=tk.TOP, fill=tk.X, padx=5, pady=5)
-        btn_select_folder = tk.Button(
-            top_frame, text="Wybierz folder", command=self.select_folder
-        )
-        btn_select_folder.pack(side=tk.LEFT, padx=5)
-        btn_prev = tk.Button(top_frame, text="<<", command=self.show_prev)
-        btn_prev.pack(side=tk.LEFT, padx=5)
-        btn_next = tk.Button(top_frame, text=">>", command=self.show_next)
-        btn_next.pack(side=tk.LEFT, padx=5)
-
         # Główna ramka – trzy kolumny
         main_frame = tk.Frame(self)
         main_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
@@ -486,15 +486,9 @@ class ImageBrowser(tk.Tk):
         manual_frame = tk.LabelFrame(right_frame, text="Manual Tags")
         manual_frame.grid(row=2, column=0, sticky=tk.NSEW, padx=5, pady=5)
         self.manual_tags_manager = TagManager(
-            manual_frame, self.settings, self.tags_repo
+            manual_frame, self.settings, self.tags_repo, self.update_final_tags
         )
         self.manual_tags_manager.pack(fill=tk.BOTH, expand=True)
-        self.manual_tags_manager.bind(
-            "<KeyRelease>", lambda e: self.update_final_tags(), add="+"
-        )
-        self.manual_tags_manager.bind(
-            "<Button-1>", lambda e: self.update_final_tags(), add="+"
-        )
 
         # Final Tags – widget Text z scrollbar, który rozszerza się
         final_frame = tk.LabelFrame(right_frame, text="Final Tags")
@@ -526,6 +520,10 @@ class ImageBrowser(tk.Tk):
         )
         self.view_post_button.pack(side=tk.LEFT, padx=5)
         self.view_post_button.config(state=tk.DISABLED)
+        btn_prev = tk.Button(upload_frame, text="<<", command=self.show_prev)
+        btn_prev.pack(side=tk.LEFT, padx=5)
+        btn_next = tk.Button(upload_frame, text=">>", command=self.show_next)
+        btn_next.pack(side=tk.LEFT, padx=5)
 
         # Na końcu okna głównego dodaj status bar:
         self.status_label = tk.Label(
@@ -601,8 +599,30 @@ class ImageBrowser(tk.Tk):
         """
         Po załadowaniu plików, sprawdza czy są jakieś pliki do uploadu oraz przetwarza Taggerem pliki.
         """
-        threading.Thread(target=self.check_uploaded_files, daemon=True).start()
-        threading.Thread(target=self.process_tagger_queue, daemon=True).start()
+        self.join_check_uploaded_files_thread()
+        self.check_uploaded_files_thread = threading.Thread(
+            target=self.check_uploaded_files
+        )
+        self.check_uploaded_files_thread.start()
+        self.join_process_tagger_queue_thread()
+        self.process_tagger_queue_thread = threading.Thread(
+            target=self.process_tagger_queue
+        )
+        self.process_tagger_queue_thread.start()
+
+    def join_check_uploaded_files_thread(self):
+        if self.check_uploaded_files_thread is not None:
+            self.check_uploaded_files_stop_event.set()
+            self.check_uploaded_files_thread.join()
+            self.check_uploaded_files_thread = None
+            self.check_uploaded_files_stop_event = threading.Event()
+
+    def join_process_tagger_queue_thread(self):
+        if self.process_tagger_queue_thread is not None:
+            self.process_tagger_queue_stop_event.set()
+            self.process_tagger_queue_thread.join()
+            self.process_tagger_queue_thread = None
+            self.process_tagger_queue_stop_event = threading.Event()
 
     def check_uploaded_files(self):
         """
@@ -617,6 +637,8 @@ class ImageBrowser(tk.Tk):
 
         batch_size = 100
         for i in range(0, len(file_md5_list), batch_size):
+            if self.check_uploaded_files_stop_event.is_set():
+                break
             batch = file_md5_list[i : i + batch_size]
             batch_md5 = [item[2] for item in batch]
             md5_param = ",".join(batch_md5)
@@ -627,11 +649,15 @@ class ImageBrowser(tk.Tk):
                 root = response.json()
                 found = {}
                 for elem in root:
+                    if self.check_uploaded_files_stop_event.is_set():
+                        break
                     post_md5 = elem.get("md5", "").lower()
                     post_id = elem.get("id")
                     if post_md5 and post_id:
                         found[post_md5] = post_id
                 for idx, file_path, md5 in batch:
+                    if self.check_uploaded_files_stop_event.is_set():
+                        break
                     self.upload_verified += 1  # Każdy plik w batchu jest zweryfikowany
                     if md5.lower() in found:
                         self.uploaded[file_path] = found[md5.lower()]
@@ -650,6 +676,7 @@ class ImageBrowser(tk.Tk):
                     self.after(0, self.update_status_bar)
             except Exception as e:
                 print("Błąd podczas sprawdzania paczki uploadu:", e)
+        self.after(100, self.join_check_uploaded_files_thread)
 
     def update_button_states(self):
         """
@@ -774,7 +801,11 @@ class ImageBrowser(tk.Tk):
             # Uaktualnij widget PNG Tags
             self.update_png_tags_widget(parsed_parameters.split())
             # Uruchom Taggera w osobnym wątku
-            threading.Thread(target=self.run_tagger, daemon=True).start()
+            thread_name = f"tagger{self.tagger_thread_idx}"
+            thread = threading.Thread(target=self.run_tagger, args=(thread_name,))
+            self.run_tagger_threads[thread_name] = thread
+            self.tagger_thread_idx += 1
+            thread.start()
         except Exception as e:
             messagebox.showerror("Błąd", f"Nie można załadować obrazka:\n{e}")
 
@@ -1196,6 +1227,8 @@ class ImageBrowser(tk.Tk):
     def process_tagger_queue(self):
         """Przetwarza wszystkie obrazki w tle (pomijając aktualnie wybrany)."""
         for file_path in self.image_files:
+            if self.process_tagger_queue_stop_event.is_set():
+                break
             # Jeśli obrazek jest aktualnie wybrany, pomijamy – on będzie przetwarzany w foreground
             if (
                 self.current_index is not None
@@ -1203,8 +1236,9 @@ class ImageBrowser(tk.Tk):
             ):
                 continue
             self.process_tagger_for_image(file_path)
+        self.after(100, self.join_process_tagger_queue_thread)
 
-    def run_tagger(self):
+    def run_tagger(self, thread_name):
         """
         Przetwarza aktualnie wybrany obrazek.
         Jeśli wynik jest już w cache, wykorzystuje go; w przeciwnym razie uruchamia Taggera
@@ -1223,9 +1257,15 @@ class ImageBrowser(tk.Tk):
         self.rating_var.set(new_rating)
         self.after(0, lambda: self.update_tagger_tags_widget(result))
         self.update_listbox_item_color_by_rating(file_path, result.rating)
+        self.after(100, lambda: self.run_tagger_threads[thread_name].join())
 
     def upload_file(
-        self, file_path, final_tags=None, final_rating=None, progress_queue=None
+        self,
+        file_path,
+        final_tags=None,
+        final_rating=None,
+        progress_queue=None,
+        cancel_event=None,
     ):
         base_file_name = os.path.basename(file_path)
         if progress_queue:
@@ -1258,7 +1298,9 @@ class ImageBrowser(tk.Tk):
                     progress_queue.put(("progress", percentage))
 
             with open(file_path, "rb") as f:
-                wrapped_file = ProgressFile(f, progress_callback, total_size)
+                wrapped_file = ProgressFile(
+                    f, progress_callback, total_size, cancel_event
+                )
                 files = {"file": (base_file_name, wrapped_file, "image/png")}
                 response = requests.post(url, data=fields, files=files)
             if progress_queue:
@@ -1293,25 +1335,21 @@ class ImageBrowser(tk.Tk):
                 self.uploaded[file_path] = post_id
                 self.uploaded_count += 1
                 self.after(0, self.update_status_bar)
-            self.after(0, self.update_button_states)
         except Exception as e:
-            self.upload_button.after(0, self.update_button_states)
             messagebox.showerror("Błąd uploadu", str(e))
+        finally:
+            self.upload_button.after(0, self.update_button_states)
 
     def edit_file(
-        self, file_path, final_tags=None, final_rating=None, progress_queue=None
+        self,
+        file_path,
+        final_tags=None,
+        final_rating=None,
+        progress_queue=None,
+        cancel_event=None,
     ):
         """
         Update tags and rating for an existing post without uploading the file.
-
-        Args:
-            file_path: Path to the image file associated with the post
-            final_tags: Optional override tags (space-separated string)
-            final_rating: Optional override rating
-            progress_queue: Progress communication queue
-
-        Requires:
-            - The file must have an existing post ID in self.uploaded
         """
         base_file_name = os.path.basename(file_path)
         post_id = self.uploaded.get(file_path)
@@ -1328,17 +1366,28 @@ class ImageBrowser(tk.Tk):
             progress_queue.put(("label", f"Aktualizuję tagi dla {base_file_name}..."))
 
         try:
+            # Check for cancellation before starting the operation.
+            if cancel_event is not None and cancel_event.is_set():
+                if progress_queue:
+                    progress_queue.put(("label", "Operacja anulowana"))
+                return
+
             # Get authentication session and token
             session = login(self.settings)
             auth_token = get_auth_token(session, self.settings)
 
+            # Check cancellation after login if needed.
+            if cancel_event is not None and cancel_event.is_set():
+                if progress_queue:
+                    progress_queue.put(("label", "Operacja anulowana"))
+                return
+
             # Prepare tags and rating
             tags = (
                 self.final_tags_text.get("1.0", tk.END).strip()
                 if final_tags is None
                 else final_tags
             )
-
             rating_value = self.rating_map.get(
                 self.rating_var.get() if final_rating is None else final_rating, "?"
             )
@@ -1358,6 +1407,12 @@ class ImageBrowser(tk.Tk):
             if progress_queue:
                 progress_queue.put(("progress", 50))
 
+            # Check for cancellation before sending the update request.
+            if cancel_event is not None and cancel_event.is_set():
+                if progress_queue:
+                    progress_queue.put(("label", "Operacja anulowana"))
+                return
+
             # Send update request
             response = session.post(url, data=payload, allow_redirects=False)
 
@@ -1365,11 +1420,9 @@ class ImageBrowser(tk.Tk):
             if response.status_code == 302:
                 if progress_queue:
                     progress_queue.put(("progress", 100))
-
                 message = "Tagi zostały zaktualizowane!"
                 if not final_tags:  # Only show success if not bulk operation
                     messagebox.showinfo("Sukces edycji", message)
-
                 # Update UI state
                 self.after(0, self.update_button_states)
                 return
@@ -1400,8 +1453,22 @@ class ImageBrowser(tk.Tk):
         ):
             return
 
-        def worker():
+        def worker(progress_queue: queue = None, cancel_event: threading.Event = None):
+            files_count = len(self.image_files)
+            if progress_queue:
+                progress_queue.put(("mode", "determinate"))
+                progress_queue.put(("max", 100))
+            file_idx = 0
             for file_path in self.image_files:
+                if progress_queue:
+                    progress_queue.put(("progress", file_idx * 1.0 / files_count))
+                    progress_queue.put(
+                        ("label", f"Wysyłam plik {file_idx+1}/{files_count}...")
+                    )
+                if cancel_event is not None and cancel_event.is_set():
+                    if progress_queue:
+                        progress_queue.put(("label", f"Anulowano operację!"))
+                    return
                 if not self.uploaded.get(file_path, False):
                     final_tags, final_rating = (
                         self.compute_final_tags_and_rating_for_file(file_path)
@@ -1412,5 +1479,8 @@ class ImageBrowser(tk.Tk):
                     self.upload_file(
                         file_path, final_tags=final_tags, final_rating=final_rating
                     )
+            if progress_queue:
+                progress_queue.put(("label", f"Przesłano pliki!"))
+                progress_queue.put(("progress", 100))
 
-        threading.Thread(target=worker, daemon=True).start()
+        self.processing_dialog = ProcessingDialog(self, worker)