fix(mastodon): improve autocompletion scan performance and UX

src/extra/autocompletionUsers/mastodon/scan.py

@@ -1,13 +1,24 @@
 # -*- coding: utf-8 -*-
 """ Scanning code for autocompletion feature on TWBlue. This module can retrieve user objects from the selected Mastodon account automatically. """
 import time
+import logging
 import wx
 import widgetUtils
 import output
+from enum import Enum
 from pubsub import pub
 from . import wx_scan
 from extra.autocompletionUsers import manage, storage
 
+log = logging.getLogger("extra.autocompletionUsers.mastodon.scan")
+
+
+class ScanType(Enum):
+    """Enum to distinguish between scanning followers or following users."""
+    FOLLOWING = "following"
+    FOLLOWERS = "followers"
+
+
 class autocompletionScan(object):
     def __init__(self, config, buffer, window):
         """ Class constructor. This class will take care of scanning the selected Mastodon account to populate the database with users automatically upon request.
@@ -23,6 +34,7 @@ def __init__(self, config, buffer, window):
         self.config = config
         self.buffer = buffer
         self.window = window
+        self.progress_dialog = None
 
     def show_dialog(self):
         """ displays a dialog to confirm which buffers should be scanned (followers or following users). """
@@ -33,67 +45,119 @@ def show_dialog(self):
             confirmation = wx_scan.confirm()
             return confirmation
 
+    def get_user_counts(self):
+        """Get the followers and following counts from the user's account."""
+        try:
+            credentials = self.buffer.session.api.account_verify_credentials()
+            followers_count = credentials.followers_count if self.dialog.get("followers") else 0
+            following_count = credentials.following_count if self.dialog.get("friends") else 0
+            return followers_count, following_count
+        except Exception as e:
+            log.exception(f"Error getting user counts: {e}")
+            return 0, 0
+
     def prepare_progress_dialog(self):
-        self.progress_dialog = wx_scan.autocompletionScanProgressDialog()
-        # connect method to update progress dialog
-        pub.subscribe(self.on_update_progress, "on-update-progress")
+        followers_count, following_count = self.get_user_counts()
+        self.progress_dialog = wx_scan.autocompletionScanProgressDialog(
+            followers_count=followers_count,
+            following_count=following_count
+        )
         self.progress_dialog.Show()
 
-    def on_update_progress(self):
-        wx.CallAfter(self.progress_dialog.progress_bar.Pulse)
+    def update_progress(self, current_users, current_page, scanning_type):
+        """Update the progress dialog from the worker thread."""
+        if self.progress_dialog:
+            wx.CallAfter(self.progress_dialog.update_progress, current_users, current_page, scanning_type)
+
+    def is_cancelled(self):
+        """Check if the user has requested cancellation."""
+        if self.progress_dialog:
+            return self.progress_dialog.cancelled
+        return False
+
+    def _scan_users(self, scan_type, users_dict):
+        """Scan users of the specified type and add them to the users dictionary.
+
+        :param scan_type: ScanType.FOLLOWING or ScanType.FOLLOWERS
+        :param users_dict: Dictionary to store users, keyed by user ID
+        :returns: True if completed, False if cancelled
+        """
+        # Select the appropriate API method
+        if scan_type == ScanType.FOLLOWING:
+            api_method = self.buffer.session.api.account_following
+        else:
+            api_method = self.buffer.session.api.account_followers
+
+        log.debug(f"Scanning {scan_type.value}...")
+        current_page = 1
+        first_page = api_method(id=self.buffer.session.db["user_id"], limit=80)
+        self.update_progress(len(users_dict), current_page, scan_type.value)
+
+        if first_page != None:
+            for user in first_page:
+                if user.id not in users_dict:
+                    users_dict[user.id] = user
+
+        next_page = first_page
+        while next_page != None:
+            if self.is_cancelled():
+                log.info(f"Scan cancelled by user during {scan_type.value} scan.")
+                return False
+            time.sleep(0.25)  # Small delay to avoid rate limiting
+            next_page = self.buffer.session.api.fetch_next(next_page)
+            current_page += 1
+            self.update_progress(len(users_dict), current_page, scan_type.value)
+            if next_page == None:
+                break
+            for user in next_page:
+                if user.id not in users_dict:
+                    users_dict[user.id] = user
+            log.debug(f"Scanned {len(users_dict)} users so far...")
+
+        return True
 
     def scan(self):
         """ Attempts to add all users selected by current user to the autocomplete database. """
         self.config["mysc"]["save_friends_in_autocompletion_db"] = self.dialog.get("friends")
         self.config["mysc"]["save_followers_in_autocompletion_db"] = self.dialog.get("followers")
-        output.speak(_("Updating database... You can close this window now. A message will tell you when the process finishes."))
+        output.speak(_("Scanning account. Please wait or press cancel to stop."))
         database = storage.storage(self.buffer.session.session_id)
-        percent = 0
-        users = []
-        if self.dialog.get("friends") == True:
-            first_page = self.buffer.session.api.account_following(id=self.buffer.session.db["user_id"], limit=80)
-            pub.sendMessage("on-update-progress")
-            if first_page != None:
-                for user in first_page:
-                    users.append(user)
-            next_page = first_page
-            while next_page != None:
-                next_page = self.buffer.session.api.fetch_next(next_page)
-                pub.sendMessage("on-update-progress")
-                if next_page == None:
-                    break
-                for user in next_page:
-                    users.append(user)
-        # same step, but for followers.
-        if self.dialog.get("followers") == True:
-            first_page = self.buffer.session.api.account_followers(id=self.buffer.session.db["user_id"], limit=80)
-            pub.sendMessage("on-update-progress")
-            if first_page != None:
-                for user in first_page:
-                    if user not in users:
-                        users.append(user)
-            next_page = first_page
-            while next_page != None:
-                next_page = self.buffer.session.api.fetch_next(next_page)
-                pub.sendMessage("on-update-progress")
-                if next_page == None:
-                    break
-                for user in next_page:
-                    if user not in users:
-                        users.append(user)
-#            except TweepyException:
-#                wx.CallAfter(wx_scan.show_error)
-#                return self.done()
-        for user in users:
-            name = user.display_name if user.display_name != None and user.display_name != "" else user.username
-            database.set_user(user.acct, name, 1)
-        wx.CallAfter(wx_scan .show_success, len(users))
+        # Use a dictionary keyed by user ID for O(1) lookups instead of O(n) list comparisons
+        users_dict = {}
+        cancelled = False
+        try:
+            if self.dialog.get("friends") == True:
+                if not self._scan_users(ScanType.FOLLOWING, users_dict):
+                    cancelled = True
+            if self.dialog.get("followers") == True and not cancelled:
+                if not self._scan_users(ScanType.FOLLOWERS, users_dict):
+                    cancelled = True
+        except Exception as e:
+            log.exception(f"Error scanning account: {e}")
+            wx.CallAfter(wx_scan.show_error)
+            return self.done()
+        # Save users to database
+        users = list(users_dict.values())
+        new_users_count = 0
+        if len(users) > 0:
+            log.debug(f"Saving {len(users)} users to autocompletion database...")
+            wx.CallAfter(self.progress_dialog.set_saving_status)
+            for user in users:
+                name = user.display_name if user.display_name != None and user.display_name != "" else user.username
+                if database.set_user(user.acct, name, 1):
+                    new_users_count += 1
+        already_existed = len(users) - new_users_count
+        if cancelled:
+            log.info(f"Scan cancelled. Found {len(users)} users, {new_users_count} new, {already_existed} already in database.")
+            wx.CallAfter(wx_scan.show_cancelled, len(users), new_users_count)
+        else:
+            log.info(f"Successfully imported {new_users_count} new users ({already_existed} already existed).")
+            wx.CallAfter(wx_scan.show_success, len(users), new_users_count)
         self.done()
 
     def done(self):
         wx.CallAfter(self.progress_dialog.Destroy)
         wx.CallAfter(self.dialog.Destroy)
-        pub.unsubscribe(self.on_update_progress, "on-update-progress")
 
 def add_user(session, database, user):
     """ Adds an user to the database. """

src/extra/autocompletionUsers/mastodon/wx_scan.py

@@ -3,6 +3,9 @@
 import widgetUtils
 import application
 
+def returnTrue():
+    return True
+
 class autocompletionScanDialog(widgetUtils.BaseDialog):
     def __init__(self):
         super(autocompletionScanDialog, self).__init__(parent=None, id=-1, title=_(u"Autocomplete users' settings"))
@@ -22,22 +25,120 @@ def __init__(self):
         self.SetClientSize(sizer.CalcMin())
 
 class autocompletionScanProgressDialog(widgetUtils.BaseDialog):
-    def __init__(self, *args, **kwargs):
+    def __init__(self, followers_count=0, following_count=0, *args, **kwargs):
         super(autocompletionScanProgressDialog, self).__init__(parent=None, id=wx.ID_ANY, title=_("Updating autocompletion database"), *args, **kwargs)
-        panel = wx.Panel(self)
+        self.cancelled = False
+        self.followers_count = followers_count
+        self.following_count = following_count
+        self.total_users = followers_count + following_count
+        panel = wx.Panel(self, style=wx.TAB_TRAVERSAL)
         sizer = wx.BoxSizer(wx.VERTICAL)
-        self.progress_bar = wx.Gauge(parent=panel)
-        sizer.Add(self.progress_bar)
-        panel.SetSizerAndFit(sizer)
+        # Account information label and field
+        info_label_text = wx.StaticText(panel, label=_("Account information:"))
+        sizer.Add(info_label_text, 0, wx.LEFT | wx.TOP, 5)
+        info_text = _("Followers: {followers} | Following: {following} | Total: {total}").format(
+            followers=followers_count,
+            following=following_count,
+            total=self.total_users
+        )
+        self.info_field = wx.TextCtrl(panel, value=info_text, style=wx.TE_READONLY | wx.TE_PROCESS_TAB)
+        self.info_field.AcceptsFocusFromKeyboard = returnTrue
+        sizer.Add(self.info_field, 0, wx.ALL | wx.EXPAND, 5)
+        # Current status label and field
+        status_label_text = wx.StaticText(panel, label=_("Current status:"))
+        sizer.Add(status_label_text, 0, wx.LEFT | wx.TOP, 5)
+        self.status_field = wx.TextCtrl(panel, value=_("Preparing..."), style=wx.TE_READONLY | wx.TE_PROCESS_TAB)
+        self.status_field.AcceptsFocusFromKeyboard = returnTrue
+        sizer.Add(self.status_field, 0, wx.ALL | wx.EXPAND, 5)
+        # Progress label and field
+        progress_label_text = wx.StaticText(panel, label=_("Progress:"))
+        sizer.Add(progress_label_text, 0, wx.LEFT | wx.TOP, 5)
+        self.progress_field = wx.TextCtrl(panel, value="", style=wx.TE_READONLY | wx.TE_PROCESS_TAB)
+        self.progress_field.AcceptsFocusFromKeyboard = returnTrue
+        sizer.Add(self.progress_field, 0, wx.ALL | wx.EXPAND, 5)
+        # Progress bar
+        self.progress_bar = wx.Gauge(parent=panel, range=self.total_users if self.total_users > 0 else 100)
+        sizer.Add(self.progress_bar, 0, wx.ALL | wx.EXPAND, 5)
+        # Cancel button
+        self.cancel_button = wx.Button(panel, wx.ID_CANCEL, _("&Cancel"))
+        self.cancel_button.Bind(wx.EVT_BUTTON, self.on_cancel)
+        sizer.Add(self.cancel_button, 0, wx.ALL | wx.ALIGN_CENTER, 10)
+        panel.SetSizer(sizer)
+        self.SetClientSize(sizer.CalcMin())
+        self.SetSize(400, -1)
+        # Handle window close - same as cancel, don't close immediately
+        self.Bind(wx.EVT_CLOSE, self.on_close)
+
+    def on_close(self, event):
+        """Handle window close event (X button). Same as cancel, but prevent immediate close."""
+        if not self.cancelled:
+            self.on_cancel()
+        # Don't call event.Skip() - this prevents the window from closing
+        # The window will be destroyed by done() after the scan thread finishes
+
+    def on_cancel(self, event=None):
+        self.cancelled = True
+        self._update_field_preserving_cursor(self.status_field, _("Cancelling... Please wait."))
+        self.cancel_button.Disable()
+
+    def _update_field_preserving_cursor(self, field, new_value):
+        """Update a text field while preserving cursor position."""
+        cursor_pos = field.GetInsertionPoint()
+        field.SetValue(new_value)
+        # Restore cursor, but don't go beyond the new text length
+        new_length = len(new_value)
+        field.SetInsertionPoint(min(cursor_pos, new_length))
+
+    def update_progress(self, current_users, current_page, scanning_type):
+        """Update the progress dialog with current scan status.
+
+        :param current_users: Number of users scanned so far
+        :param current_page: Current page being processed
+        :param scanning_type: 'following' or 'followers'
+        """
+        if scanning_type == "following":
+            status = _("Scanning following users...")
+        else:
+            status = _("Scanning followers...")
+        self._update_field_preserving_cursor(self.status_field, status)
+        progress_text = _("Page {page} | Users processed: {current} / {total}").format(
+            page=current_page,
+            current=current_users,
+            total=self.total_users
+        )
+        self._update_field_preserving_cursor(self.progress_field, progress_text)
+        # Update progress bar
+        if self.total_users > 0:
+            self.progress_bar.SetValue(min(current_users, self.total_users))
+
+    def set_saving_status(self):
+        """Update status to show we're saving to database."""
+        self._update_field_preserving_cursor(self.status_field, _("Saving users to database..."))
+        self._update_field_preserving_cursor(self.progress_field, "")
+        self.progress_bar.SetValue(self.total_users)
 
 def confirm():
     with wx.MessageDialog(None, _("This process will retrieve the users you selected from your Mastodon account, and add them to the user autocomplete database. Please note that if there are many users or you have tried to perform this action less than 15 minutes ago, TWBlue may reach a limit in API calls when trying to load the users into the database. If this happens, we will show you an error, in which case you will have to try this process again in a few minutes. If this process ends with no error, you will be redirected back to the account settings dialog. Do you want to continue?"), _("Attention"), style=wx.ICON_QUESTION|wx.YES_NO) as result:
         if result.ShowModal() == wx.ID_YES:
             return True
         return False
 
-def show_success(users):
-    with wx.MessageDialog(None, _("TWBlue has imported {} users successfully.").format(users), _("Done")) as dlg:
+def show_success(total_users, new_users):
+    already_existed = total_users - new_users
+    message = _("Scan completed. {new} new users imported, {existing} already in database.").format(
+        new=new_users,
+        existing=already_existed
+    )
+    with wx.MessageDialog(None, message, _("Done")) as dlg:
+        dlg.ShowModal()
+
+def show_cancelled(total_users, new_users):
+    already_existed = total_users - new_users
+    message = _("Operation cancelled. {new} new users imported, {existing} already in database.").format(
+        new=new_users,
+        existing=already_existed
+    )
+    with wx.MessageDialog(None, message, _("Cancelled"), style=wx.ICON_INFORMATION) as dlg:
         dlg.ShowModal()
 
 def show_error():

src/extra/autocompletionUsers/storage.py

@@ -25,8 +25,13 @@ def get_users(self, term):
         return self.cursor.fetchall()
 
     def set_user(self, screen_name, user_name, from_a_buffer):
+        """Insert a user into the database.
+
+        :returns: True if the user was inserted (new), False if already existed.
+        """
         self.cursor.execute("""insert or ignore into users values(?, ?, ?)""", (screen_name, user_name, from_a_buffer))
         self.connection.commit()
+        return self.cursor.rowcount > 0
 
     def remove_user(self, user):
         self.cursor.execute("""DELETE FROM users WHERE user = ?""", (user,))
@@ -48,5 +53,8 @@ def create_table(self):
 )""")
 
     def __del__(self):
-        self.cursor.close()
-        self.connection.close()
+        try:
+            self.cursor.close()
+            self.connection.close()
+        except Exception:
+            pass  # Already closed

src/test/extra/__init__.py

@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-

src/test/extra/autocompletionUsers/__init__.py

@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-