From 08a9bf5667674b645189ee95c43dbcf1d1cdcc7e Mon Sep 17 00:00:00 2001
From: Frank LENORMAND <lenormf@gmail.com>
Date: Fri, 1 May 2020 13:50:27 +0300
Subject: [PATCH] Don't lose the focus when refreshing a thread

This commit allows keeping the currently focused message selected
across thread refreshes.

Since this operation can take a long time in large threads, a
`search_threads_rebuild_limit` option is introduced. When defined,
it sets the maximum amount of messages that will be iterated upon
when trying to find the previously focused message.
---
 alot/buffers/search.py                 | 28 ++++++++++++++++++++++++++
 alot/defaults/alot.rc.spec             |  4 ++++
 docs/source/configuration/alotrc_table | 11 ++++++++++
 3 files changed, 43 insertions(+)

diff --git a/alot/buffers/search.py b/alot/buffers/search.py
index d64e857f3..1bd649f1e 100644
--- a/alot/buffers/search.py
+++ b/alot/buffers/search.py
@@ -25,7 +25,9 @@ def __init__(self, ui, initialquery='', sort_order=None):
         default_order = settings.get('search_threads_sort_order')
         self.sort_order = sort_order or default_order
         self.result_count = 0
+        self.search_threads_rebuild_limit = settings.get('search_threads_rebuild_limit')
         self.isinitialized = False
+        self.threadlist = None
         self.rebuild()
         Buffer.__init__(self, ui, self.body)
 
@@ -50,6 +52,11 @@ def rebuild(self, reverse=False):
         else:
             order = self.sort_order
 
+        if self.threadlist:
+            selected_thread = self.get_selected_thread()
+        else:
+            selected_thread = None
+
         exclude_tags = settings.get_notmuch_setting('search', 'exclude_tags')
         if exclude_tags:
             exclude_tags = [t for t in exclude_tags.split(';') if t]
@@ -72,6 +79,9 @@ def rebuild(self, reverse=False):
         self.listbox = urwid.ListBox(self.threadlist)
         self.body = self.listbox
 
+        if selected_thread:
+            self.focus_thread(selected_thread)
+
     def get_selected_threadline(self):
         """
         returns curently focussed :class:`alot.widgets.ThreadlineWidget`
@@ -92,6 +102,14 @@ def consume_pipe(self):
         while not self.threadlist.empty:
             self.threadlist._get_next_item()
 
+    def consume_pipe_until(self, predicate, limit=0):
+        n = limit
+        while not limit or n > 0:
+            if self.threadlist.empty \
+               or predicate(self.threadlist._get_next_item()):
+                break
+            n -= 1
+
     def focus_first(self):
         if not self.reversed:
             self.body.set_focus(0)
@@ -108,3 +126,13 @@ def focus_last(self):
         else:
             self.rebuild(reverse=True)
 
+    def focus_thread(self, thread):
+        tid = thread.get_thread_id()
+        self.consume_pipe_until(lambda w:
+                                w and w.get_thread().get_thread_id() == tid,
+                                self.search_threads_rebuild_limit)
+
+        for pos, threadlinewidget in enumerate(self.threadlist.get_lines()):
+            if threadlinewidget.get_thread().get_thread_id() == tid:
+                self.body.set_focus(pos)
+                break
diff --git a/alot/defaults/alot.rc.spec b/alot/defaults/alot.rc.spec
index 18d907927..c64789b18 100644
--- a/alot/defaults/alot.rc.spec
+++ b/alot/defaults/alot.rc.spec
@@ -188,6 +188,10 @@ initial_command = string(default='search tag:inbox AND NOT tag:killed')
 # default sort order of results in a search
 search_threads_sort_order = option('oldest_first', 'newest_first', 'message_id', 'unsorted', default='newest_first')
 
+# maximum amount of threads that will be consumed to try to restore the focus, upon triggering a search buffer rebuild
+# when set to 0, no limit is set (can be very slow in searches that yield thousands of results)
+search_threads_rebuild_limit = integer(default=0)
+
 # in case more than one account has an address book:
 # Set this to True to make tab completion for recipients during compose only
 # look in the abook of the account matching the sender address
diff --git a/docs/source/configuration/alotrc_table b/docs/source/configuration/alotrc_table
index 5e129e778..3ee58ccc0 100644
--- a/docs/source/configuration/alotrc_table
+++ b/docs/source/configuration/alotrc_table
@@ -598,6 +598,17 @@
     :default: newest_first
 
 
+.. _search-threads-rebuild-limit:
+
+.. describe:: search_threads_rebuild_limit
+
+    maximum amount of threads that will be consumed to try to restore the focus, upon triggering a search buffer rebuild
+    when set to 0, no limit is set (can be very slow in searches that yield thousands of results)
+
+    :type: integer
+    :default: 0
+
+
 .. _show-statusbar:
 
 .. describe:: show_statusbar