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.

7 files changed, 64 insertions, 72 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
diff --git a/tests/db/test_message.py b/tests/db/test_message.py
index 953d0fe4..43b06e19 100644
--- a/tests/db/test_message.py
+++ b/tests/db/test_message.py
@@ -66,7 +66,7 @@ class TestMessage(unittest.TestCase):
         """
         msg = message.Message(mock.Mock(),
                               MockNotmuchMessage({'From': 'user@example.com'}),
-                              mock.Mock())
+                              mock.Mock(), mock.Mock())
         self.assertEqual(msg.get_author(), ('', 'user@example.com'))
 
     def test_get_author_name_and_email(self):
@@ -76,7 +76,7 @@ class TestMessage(unittest.TestCase):
         msg = message.Message(
             mock.Mock(),
             MockNotmuchMessage({'From': '"User Name" <user@example.com>'}),
-            mock.Mock())
+            mock.Mock(), mock.Mock())
         self.assertEqual(msg.get_author(), ('User Name', 'user@example.com'))
 
     def test_get_author_sender(self):
@@ -86,7 +86,7 @@ class TestMessage(unittest.TestCase):
         msg = message.Message(
             mock.Mock(),
             MockNotmuchMessage({'Sender': '"User Name" <user@example.com>'}),
-            mock.Mock())
+            mock.Mock(), mock.Mock())
         self.assertEqual(msg.get_author(), ('User Name', 'user@example.com'))
 
     def test_get_author_no_name_draft(self):
@@ -99,7 +99,7 @@ class TestMessage(unittest.TestCase):
         with mock.patch('alot.db.message.settings.get_accounts',
                         mock.Mock(return_value=[acc])):
             msg = message.Message(
-                mock.Mock(), MockNotmuchMessage(tags=['draft']), mock.Mock())
+                mock.Mock(), MockNotmuchMessage(tags=['draft']), mock.Mock(), mock.Mock())
         self.assertEqual(msg.get_author(), ('User Name', 'user@example.com'))
 
     def test_get_author_no_name(self):
@@ -111,5 +111,5 @@ class TestMessage(unittest.TestCase):
         acc.realname = 'User Name'
         with mock.patch('alot.db.message.settings.get_accounts',
                         mock.Mock(return_value=[acc])):
-            msg = message.Message(mock.Mock(), MockNotmuchMessage(), mock.Mock())
+            msg = message.Message(mock.Mock(), MockNotmuchMessage(), mock.Mock(), mock.Mock())
         self.assertEqual(msg.get_author(), ('Unknown', ''))
diff --git a/tests/db/test_thread.py b/tests/db/test_thread.py
index 8c78a026..8b43b0b9 100644
--- a/tests/db/test_thread.py
+++ b/tests/db/test_thread.py
@@ -38,15 +38,15 @@ class TestThreadGetAuthor(unittest.TestCase):
                                                minute=10)),
                      ('ooh', None)]:
             m = mock.Mock()
-            m.get_date = mock.Mock(return_value=d)
+            m.date = d
             m.get_author = mock.Mock(return_value=a)
             get_messages.append(m)
         gm = mock.Mock()
-        gm.keys = mock.Mock(return_value=get_messages)
+        gm.values = mock.Mock(return_value=get_messages)
 
         cls.__patchers.extend([
-            mock.patch('alot.db.thread.Thread.get_messages',
-                       new=mock.Mock(return_value=gm)),
+            mock.patch('alot.db.thread.Thread.messages',
+                       new=mock.Mock(return_value=get_messages)),
             mock.patch('alot.db.thread.Thread.refresh', new=mock.Mock()),
         ])