Fix /correct and /me highlights, and handle /correct a bit better.

4 files changed, 48 insertions, 39 deletions

diff --git a/src/core.py b/src/core.py
index f63ff062..dd8e26af 100644
--- a/src/core.py
+++ b/src/core.py
@@ -2455,7 +2455,8 @@ class Core(object):
             date = date if delayed == True else None
             replaced_id = message['replace']['id']
             if replaced_id is not '':
-                tab.modify_message(body, replaced_id, message['id'])
+                if tab.modify_message(body, replaced_id, message['id'], date, nick_from):
+                    self.events.trigger('highlight', message, tab)
             elif tab.add_message(body, date, nick_from, history=True if date else False, identifier=message['id']):
                 self.events.trigger('highlight', message, tab)
             if tab is self.current_tab():
diff --git a/src/tabs.py b/src/tabs.py
index 037d591f..d9dae765 100644
--- a/src/tabs.py
+++ b/src/tabs.py
@@ -590,11 +590,6 @@ class ChatTab(Tab):
     def completion_correct(self, the_input):
         return the_input.auto_completion([self.last_sent_message['body']], '', quotify=False)
 
-    def modify_message(self, txt, old_id, new_id):
-        self._text_buffer.modify_message(txt, old_id, new_id)
-        self.text_win.rebuild_everything(self._text_buffer)
-        self.core.refresh_window()
-
     @property
     def inactive(self):
         """Whether we should send inactive or active as a chatstate"""
@@ -1641,6 +1636,14 @@ class MucTab(ChatTab):
         self._text_buffer.add_message(txt, time, nickname, nick_color, history, user, highlight=highlight, identifier=identifier)
         return highlight
 
+    def modify_message(self, txt, old_id, new_id, time=None, nickname=None):
+        self.log_message(txt, time, nickname)
+        highlight = self.do_highlight(txt, time, nickname)
+        self._text_buffer.modify_message(txt, old_id, new_id, highlight=highlight, time=time)
+        self.text_win.rebuild_everything(self._text_buffer)
+        self.core.refresh_window()
+        return highlight
+
     def matching_names(self):
         return [safeJID(self.get_name()).user]
 
diff --git a/src/text_buffer.py b/src/text_buffer.py
index 2d3cb87a..67aa8ef6 100644
--- a/src/text_buffer.py
+++ b/src/text_buffer.py
@@ -18,7 +18,7 @@ from datetime import datetime
 from config import config
 from theming import get_theme
 
-Message = collections.namedtuple('Message', 'txt nick_color time str_time nickname user identifier highlight')
+Message = collections.namedtuple('Message', 'txt nick_color time str_time nickname user identifier highlight me old_message')
 
 class TextBuffer(object):
     """
@@ -40,29 +40,28 @@ class TextBuffer(object):
         return self.messages[-1] if self.messages else None
 
 
-    def make_message(self, txt, time, nickname, nick_color, history, user, identifier, str_time=None, highlight=False):
+    def make_message(self, txt, time, nickname, nick_color, history, user, identifier, str_time=None, highlight=False, old_message=None):
         time = time or datetime.now()
+        me = False
         if txt.startswith('/me '):
-            if nick_color:
-                color = nick_color[0]
-            elif user:
-                color = user.color[0]
-            else:
-                color = None
-            # TODO: display the bg color too.
-            txt = '\x19%(info_col)s}* \x19%(col)s}%(nick)s \x19%(info_col)s}%(msg)s' % {'info_col':get_theme().COLOR_ME_MESSAGE[0], 'col': color or 5, 'nick': nickname, 'msg': txt[4:]}
-            nickname = None
+            me = True
+            txt = '\x19%(info_col)s}' % {'info_col': get_theme().COLOR_ME_MESSAGE[0]} + txt[4:]
         msg = Message(
                 txt='%s\x19o'%(txt.replace('\t', '    '),),
                 nick_color=nick_color,
                 time=time,
                 str_time=(time.strftime("%Y-%m-%d %H:%M:%S") if history else time.strftime("%H:%M:%S")) if str_time is None else '',
-                      nickname=nickname, user=user, identifier=identifier, highlight=highlight)
+                nickname=nickname,
+                user=user,
+                identifier=identifier,
+                highlight=highlight,
+                me=me,
+                old_message=old_message)
         log.debug('Set message %s with %s.' % (identifier, msg))
         return msg
 
     def add_message(self, txt, time=None, nickname=None, nick_color=None, history=None, user=None, highlight=False, identifier=None, str_time=None):
-        msg = self.make_message(txt, time, nickname, nick_color, history, user, identifier, str_time, highlight)
+        msg = self.make_message(txt, time, nickname, nick_color, history, user, identifier, str_time=str_time, highlight=highlight)
         self.messages.append(msg)
         while len(self.messages) > self.messages_nb_limit:
             self.messages.pop(0)
@@ -76,14 +75,14 @@ class TextBuffer(object):
                 window.scroll_up(nb)
         return ret_val or 1
 
-    def modify_message(self, txt, old_id, new_id):
+    def modify_message(self, txt, old_id, new_id, highlight=False, time=None):
         for i, msg in enumerate(self.messages):
             if msg.identifier == old_id:
-                message = self.make_message(txt, msg.time, msg.nickname, msg.nick_color, None, msg.user, new_id, msg.highlight)
+                message = self.make_message(txt, time if time else msg.time, msg.nickname, msg.nick_color, None, msg.user, new_id, highlight=highlight, old_message=msg)
                 self.messages[i] = message
                 log.debug('Replacing message %s with %s.' % (old_id, new_id))
                 return
-        log.debug('Message %s not found in text_buffer, abort replacement.' % (identifier))
+        log.debug('Message %s not found in text_buffer, abort replacement.' % (old_id))
 
     def del_window(self, win):
         self.windows.remove(win)
diff --git a/src/windows.py b/src/windows.py
index bc697300..b18102bc 100644
--- a/src/windows.py
+++ b/src/windows.py
@@ -789,7 +789,10 @@ class TextWin(Win):
         nick = truncate_nick(message.nickname)
         offset = 0
         if nick:
-            offset += wcwidth.wcswidth(nick) + 2 # + nick + spaces length
+            # TODO: add the number of corrections.
+            offset += wcwidth.wcswidth(nick) + 2 # + nick + '> ' length
+        if message.me:
+            offset += 1 # '* ' before and ' ' after
         if timestamp:
             if message.str_time:
                 offset += 1 + len(message.str_time)
@@ -804,9 +807,8 @@ class TextWin(Win):
                                              start_pos=line[0],
                                              end_pos=line[1]))
         else:
-
             for line in lines:
-                saved_line = Line(msg=message,start_pos=line[0],end_pos=line[1])
+                saved_line = Line(msg=message, start_pos=line[0], end_pos=line[1])
                 self.built_lines.append(saved_line)
                 if highlight:
                     highlight = False
@@ -817,7 +819,7 @@ class TextWin(Win):
             return len(lines)
 
     def refresh(self):
-        log.debug('Refresh: %s',self.__class__.__name__)
+        log.debug('Refresh: %s', self.__class__.__name__)
         if self.height <= 0:
             return
         if self.pos == 0:
@@ -840,18 +842,14 @@ class TextWin(Win):
                             color = None
                         if with_timestamps:
                             self.write_time(msg.str_time)
-                        if msg.highlight:
-                            hl_color = get_theme().COLOR_HIGHLIGHT_NICK
-                            if hl_color == "reverse":
-                                self._win.attron(curses.A_REVERSE)
-                            else:
-                                color = hl_color
-                            self.write_nickname(msg.nickname, color)
-                            if hl_color == "reverse":
-                                self._win.attroff(curses.A_REVERSE)
+                        if msg.me:
+                            self._win.attron(to_curses_attr(get_theme().COLOR_ME_MESSAGE))
+                            self.addstr('* ')
+                            self.write_nickname(msg.nickname, color, msg.highlight)
+                            self.addstr(' ')
                         else:
-                            self.write_nickname(msg.nickname, color)
-                        self.addstr("> ")
+                            self.write_nickname(msg.nickname, color, msg.highlight)
+                            self.addstr('> ')
                 if y != self.height-1:
                     self.addstr('\n')
             self._win.attrset(0)
@@ -863,7 +861,7 @@ class TextWin(Win):
                                     # Offset for the timestamp (if any) plus a space after it
                                     (0 if not with_timestamps else (len(line.msg.str_time) + 1)) +
                                     # Offset for the nickname (if any) plus a space and a > after it
-                                    (0 if not line.msg.nickname else (len(truncate_nick(line.msg.nickname))) + 2),
+                                    (0 if not line.msg.nickname else (len(truncate_nick(line.msg.nickname)) + (3 if line.msg.me else 2))),
                             line.msg.txt[line.start_pos:line.end_pos])
                 if y != self.height-1:
                     self.addstr('\n')
@@ -879,18 +877,26 @@ class TextWin(Win):
         """
         self.addstr_colored(txt, y, x)
 
-    def write_nickname(self, nickname, color):
+    def write_nickname(self, nickname, color, highlight=False):
         """
         Write the nickname, using the user's color
         and return the number of written characters
         """
         if not nickname:
             return
+        if highlight:
+            hl_color = get_theme().COLOR_HIGHLIGHT_NICK
+            if hl_color == "reverse":
+                self._win.attron(curses.A_REVERSE)
+            else:
+                color = hl_color
         if color:
             self._win.attron(to_curses_attr(color))
         self.addstr(truncate_nick(nickname))
         if color:
             self._win.attroff(to_curses_attr(color))
+        if highlight and hl_color == "reverse":
+            self._win.attroff(curses.A_REVERSE)
 
     def write_time(self, time):
         """