db: rewrite the API for representing message trees

It should be cleaner and easier to use, and eventually replace the
custom tree walker in the thread display buffer.

5 files changed, 55 insertions, 63 deletions

diff --git a/alot/commands/thread.py b/alot/commands/thread.py
index 9c79840f..f720e297 100644
--- a/alot/commands/thread.py
+++ b/alot/commands/thread.py
@@ -676,7 +676,7 @@ class PipeCommand(Command):
             thread = ui.current_buffer.thread
             if not thread:
                 return
-            to_print = thread.get_messages().keys()
+            to_print = thread.messages.values()
         else:
             to_print = [ui.current_buffer.get_selected_message()]
 
diff --git a/alot/db/message.py b/alot/db/message.py
index 0b59090d..f9e8e37b 100644
--- a/alot/db/message.py
+++ b/alot/db/message.py
@@ -35,7 +35,10 @@ class Message:
     """value of the Message-Id header (str)"""
     id = None
 
-    def __init__(self, dbman, msg, thread):
+    """A list of replies to this message"""
+    replies = None
+
+    def __init__(self, dbman, msg, thread, replies):
         """
         :param dbman: db manager that is used for further lookups
         :type dbman: alot.db.DBManager
@@ -43,10 +46,13 @@ class Message:
         :type msg: notmuch.database.Message
         :param thread: this messages thread
         :type thread: :class:`~alot.db.Thread`
+        :param replies: a list of replies to this message
+        :type replies alot.db.message.Message
         """
         self._dbman = dbman
         self.id = msg.get_message_id()
         self.thread = thread
+        self.replies = replies
         try:
             self.date = datetime.fromtimestamp(msg.get_date())
         except ValueError:
diff --git a/alot/db/thread.py b/alot/db/thread.py
index bae48b1a..2792737c 100644
--- a/alot/db/thread.py
+++ b/alot/db/thread.py
@@ -36,6 +36,15 @@ class Thread:
     """Thread subject"""
     subject = None
 
+    """A list of toplevel messages"""
+    toplevel_messages = None
+
+    """A list of ids of all messages in this thread in depth-first order"""
+    message_ids = None
+
+    """A dict mapping Message-Id strings to Message instances"""
+    messages = None
+
     def __init__(self, dbman, thread):
         """
         :param dbman: db manager that is used for further lookups
@@ -46,9 +55,12 @@ class Thread:
         self._dbman = dbman
         self._authors = None
         self.id = thread.get_thread_id()
-        self._messages = {}
         self._tags = set()
 
+        self.toplevel_messages = []
+        self.message_ids       = []
+        self.messages          = {}
+
         self.refresh(thread)
 
     def refresh(self, thread=None):
@@ -84,8 +96,34 @@ class Thread:
             self.newest_date = None
 
         self._tags = {t for t in thread.get_tags()}
-        self._messages = {}  # this maps messages to its children
-        self._toplevel_messages = []
+
+        self.messages, self.toplevel_messages, self.message_ids = self._gather_messages()
+
+    def _gather_messages(self):
+        query = self._dbman.query('thread:' + self.id)
+        nm_thread = next(query.search_threads())
+
+        msgs     = {}
+        msg_tree = []
+        ids      = []
+
+        def thread_tree_walk(nm_msg):
+            msg_id = nm_msg.get_message_id()
+            ids.append(msg_id)
+
+            replies = []
+            for m in nm_msg.get_replies():
+                replies.append(thread_tree_walk(m))
+            msg = Message(self._dbman, nm_msg, self, replies)
+
+            msgs[msg_id] = msg
+
+            return msg
+
+        for m in nm_thread.get_toplevel_messages():
+            msg_tree.append(thread_tree_walk(m))
+
+        return msgs, msg_tree, ids
 
     def __str__(self):
         return "thread:%s: %s" % (self.id, self.subject)
@@ -101,7 +139,7 @@ class Thread:
         """
         tags = set(list(self._tags))
         if intersection:
-            for m in self.get_messages().keys():
+            for m in self.messages.values():
                 tags = tags.intersection(set(m.get_tags()))
         return tags
 
@@ -172,7 +210,7 @@ class Thread:
         if self._authors is None:
             # Sort messages with date first (by date ascending), and those
             # without a date last.
-            msgs = sorted(self.get_messages().keys(),
+            msgs = sorted(self.messages.values(),
                           key=lambda m: m.date or datetime.max)
 
             orderby = settings.get('thread_authors_order_by')
@@ -222,58 +260,6 @@ class Thread:
         else:
             return self._notmuch_authors_string
 
-    def get_toplevel_messages(self):
-        """
-        returns all toplevel messages contained in this thread.
-        This are all the messages without a parent message
-        (identified by 'in-reply-to' or 'references' header.
-
-        :rtype: list of :class:`~alot.db.message.Message`
-        """
-        if not self._messages:
-            self.get_messages()
-        return self._toplevel_messages
-
-    def get_messages(self):
-        """
-        returns all messages in this thread as dict mapping all contained
-        messages to their direct responses.
-
-        :rtype: dict mapping :class:`~alot.db.message.Message` to a list of
-                :class:`~alot.db.message.Message`.
-        """
-        if not self._messages:  # if not already cached
-            query = self._dbman.query('thread:' + self.id)
-            thread = next(query.search_threads())
-
-            def accumulate(acc, msg):
-                M = Message(self._dbman, msg, thread=self)
-                acc[M] = []
-                r = msg.get_replies()
-                if r is not None:
-                    for m in r:
-                        acc[M].append(accumulate(acc, m))
-                return M
-
-            self._messages = {}
-            for m in thread.get_toplevel_messages():
-                self._toplevel_messages.append(accumulate(self._messages, m))
-        return self._messages
-
-    def get_replies_to(self, msg):
-        """
-        returns all replies to the given message contained in this thread.
-
-        :param msg: parent message to look up
-        :type msg: :class:`~alot.db.message.Message`
-        :returns: list of :class:`~alot.db.message.Message` or `None`
-        """
-        msg_hash = self.get_messages()
-        for m in msg_hash.keys():
-            if m.id == msg.id:
-                return msg_hash[m]
-        return None
-
     def matches(self, query):
         """
         Check if this thread matches the given notmuch query.
diff --git a/alot/widgets/search.py b/alot/widgets/search.py
index a4d22e82..4e612e8c 100644
--- a/alot/widgets/search.py
+++ b/alot/widgets/search.py
@@ -182,7 +182,7 @@ def prepare_subject_string(thread):
 
 
 def prepare_content_string(thread):
-    msgs = sorted(thread.get_messages().keys(),
+    msgs = sorted(thread.messages.values(),
                   key=lambda msg: msg.date, reverse=True)
     lastcontent = ' '.join(m.get_body_text() for m in msgs)
     lastcontent = lastcontent.replace('^>.*$', '')
diff --git a/alot/widgets/thread.py b/alot/widgets/thread.py
index 47034bfe..97a34417 100644
--- a/alot/widgets/thread.py
+++ b/alot/widgets/thread.py
@@ -332,7 +332,7 @@ class ThreadTree(Tree):
     """
     def __init__(self, thread):
         self._thread = thread
-        self.root = thread.get_toplevel_messages()[0].id
+        self.root = thread.toplevel_messages[0].id
         self._parent_of = {}
         self._first_child_of = {}
         self._last_child_of = {}
@@ -347,7 +347,7 @@ class ThreadTree(Tree):
             odd = not odd
             last = None
             self._first_child_of[mid] = None
-            for reply in thread.get_replies_to(msg):
+            for reply in msg.replies:
                 rid = reply.id
                 if self._first_child_of[mid] is None:
                     self._first_child_of[mid] = rid
@@ -360,7 +360,7 @@ class ThreadTree(Tree):
             return odd
 
         last = None
-        for msg in thread.get_toplevel_messages():
+        for msg in thread.toplevel_messages:
             mid = msg.id
             self._prev_sibling_of[mid] = last
             self._next_sibling_of[last] = mid