fix ThreadTree

1 files changed, 7 insertions, 9 deletions

diff --git a/alot/walker.py b/alot/walker.py
index c8a28506..6f11a8b3 100644
--- a/alot/walker.py
+++ b/alot/walker.py
@@ -4,6 +4,7 @@
 import urwid
 import logging
 from alot.foreign.urwidtrees import Tree
+from alot.widgets.thread import MessageWidget
 
 
 class PipeWalker(urwid.ListWalker):
@@ -93,14 +94,12 @@ class ThreadTree(Tree):
 
         for parent, childlist in thread.get_messages().items():
             pid = parent.get_message_id()
-            self._message[pid] = parent
+            self._message[pid] = MessageWidget(parent)
 
             if childlist:
                 cid = childlist[0].get_message_id()
                 self._first_child_of[pid] = cid
                 self._parent_of[cid] = pid
-                self._prev_sibling_of[cid] = None
-                self._next_sibling_of[cid] = None
 
                 last_cid = cid
                 for child in childlist[1:]:
@@ -108,19 +107,18 @@ class ThreadTree(Tree):
                     self._parent_of[cid] = pid
                     self._prev_sibling_of[cid] = last_cid
                     self._next_sibling_of[last_cid] = cid
-                self._next_sibling_of[cid] = None
 
     def parent_position(self, pos):
-        return self._parent_of[pos]
+        return self._parent_of.get(pos, None)
 
     def first_child_position(self, pos):
-        return self._first_child_of[pos]
+        return self._first_child_of.get(pos, None)
 
     def last_child_position(self, pos):
-        return self._last_child_of[pos]
+        return self._last_child_of.get(pos, None)
 
     def next_sibling_position(self, pos):
-        return self._next_sibling_of(pos)
+        return self._next_sibling_of.get(pos, None)
 
     def prev_sibling_position(self, pos):
-        return self._prev_sibling_of(pos)
+        return self._prev_sibling_of.get(pos, None)