reimplement thread.TagCommand

3 files changed, 22 insertions, 12 deletions

diff --git a/alot/buffers.py b/alot/buffers.py
index 635cc8e2..bb502ab4 100644
--- a/alot/buffers.py
+++ b/alot/buffers.py
@@ -381,7 +381,7 @@ class ThreadBuffer(Buffer):
 
 
     def refresh(self):
-        """refresh and flushe caches or Thread tree"""
+        """refresh and flushe caches of Thread tree"""
         self.body.refresh()
 
     # needed for ui.get_deep_focus..
diff --git a/alot/commands/thread.py b/alot/commands/thread.py
index 8dbdbdd3..b361b7bf 100644
--- a/alot/commands/thread.py
+++ b/alot/commands/thread.py
@@ -945,21 +945,21 @@ class TagCommand(Command):
         Command.__init__(self, **kwargs)
 
     def apply(self, ui):
-        all_message_widgets = ui.current_buffer.get_messagewidgets()
+        tbuffer = ui.current_buffer
         if self.all:
-            mwidgets = all_message_widgets
+            messagetrees = tbuffer.messagetrees()
         else:
-            mwidgets = [ui.current_buffer.get_selection()]
-        messages = [mw.get_message() for mw in mwidgets]
-        logging.debug('TAG %s' % str(messages))
+            messagetrees = [tbuffer.get_selected_messagetree()]
 
         def refresh_widgets():
-            for mw in all_message_widgets:
-                mw.rebuild()
+            for mt in messagetrees:
+                mt.refresh()
+            tbuffer.refresh()
 
         tags = filter(lambda x: x, self.tagsstring.split(','))
         try:
-            for m in messages:
+            for mt in messagetrees:
+                m = mt.get_message()
                 if self.action == 'add':
                     m.add_tags(tags, afterwards=refresh_widgets)
                 if self.action == 'set':
@@ -977,6 +977,7 @@ class TagCommand(Command):
                             to_add.append(t)
                     m.remove_tags(to_remove)
                     m.add_tags(to_add, afterwards=refresh_widgets)
+
         except DatabaseROError:
             ui.notify('index in read-only mode', priority='error')
             return
diff --git a/alot/widgets/thread.py b/alot/widgets/thread.py
index 2be80ea5..5afc0369 100644
--- a/alot/widgets/thread.py
+++ b/alot/widgets/thread.py
@@ -139,9 +139,9 @@ class DictList(SimpleTree):
 class MessageTree(CollapsibleTree):
     def __init__(self, message, odd=True):
         self._message = message
-        self._summaryw = MessageSummaryWidget(message, even=(not odd))
-
+        self._odd = odd
         self.display_source = False
+        self._summaryw = None
         self._bodytree = None
         self._sourcetree = None
         self.display_all_headers = False
@@ -158,6 +158,10 @@ class MessageTree(CollapsibleTree):
     def reassemble(self):
         self._maintree._treelist = self._assemble_structure()
 
+    def refresh(self):
+        self._summaryw = None
+        self.reassemble()
+
     def debug(self):
         logging.debug('collapsed %s' % self.is_collapsed(self.root))
         logging.debug('display_source %s' % self.display_source)
@@ -183,7 +187,7 @@ class MessageTree(CollapsibleTree):
                 mainstruct.append((self._get_body(), None))
 
         structure = [
-            (self._summaryw, mainstruct)
+            (self._get_summary(), mainstruct)
         ]
         return structure
 
@@ -191,6 +195,11 @@ class MessageTree(CollapsibleTree):
         self.set_position_collapsed(
             self.root, self._message.matches(querystring))
 
+    def _get_summary(self):
+        if self._summaryw is None:
+            self._summaryw = MessageSummaryWidget(self._message, even=(not self._odd))
+        return self._summaryw
+
     def _get_source(self):
         if self._sourcetree is None:
             sourcetxt = self._message.get_email().as_string()