Typing improvements

3 files changed, 40 insertions, 16 deletions

diff --git a/poezio/text_buffer.py b/poezio/text_buffer.py
index 3c8d78ba..2ae51d80 100644
--- a/poezio/text_buffer.py
+++ b/poezio/text_buffer.py
@@ -11,11 +11,20 @@ independently by their TextWins.
 import logging
 log = logging.getLogger(__name__)
 
-from typing import Dict, Union, Optional, List, Tuple
+from typing import (
+    Dict,
+    List,
+    Optional,
+    TYPE_CHECKING,
+    Tuple,
+    Union,
+)
 from datetime import datetime
 from poezio.config import config
 from poezio.ui.types import Message, BaseMessage
 
+if TYPE_CHECKING:
+    from poezio.windows.text_win import TextWin
 
 
 class CorrectionError(Exception):
@@ -44,13 +53,13 @@ class TextBuffer:
         # we keep track of one or more windows
         # so we can pass the new messages to them, as they are added, so
         # they (the windows) can build the lines from the new message
-        self._windows = []
+        self._windows = []  # type: List[TextWin]
 
     def add_window(self, win) -> None:
         self._windows.append(win)
 
     @property
-    def last_message(self) -> Optional[Message]:
+    def last_message(self) -> Optional[BaseMessage]:
         return self.messages[-1] if self.messages else None
 
     def add_message(self, msg: BaseMessage):
@@ -113,6 +122,8 @@ class TextBuffer:
         if i == -1:
             return None
         msg = self.messages[i]
+        if not isinstance(msg, Message):
+            return None
         if msg.ack == 1:  # Message was already acked
             return False
         if msg.jid != jid:
@@ -150,7 +161,8 @@ class TextBuffer:
             raise CorrectionError("nothing to replace")
 
         msg = self.messages[i]
-
+        if not isinstance(msg, Message):
+            raise CorrectionError('Wrong message type')
         if msg.user and msg.user is not user:
             raise CorrectionError("Different users")
         elif msg.history:
diff --git a/poezio/ui/render.py b/poezio/ui/render.py
index 3f94ecd8..b8368312 100644
--- a/poezio/ui/render.py
+++ b/poezio/ui/render.py
@@ -3,14 +3,17 @@ import curses
 
 from datetime import datetime
 from functools import singledispatch
-from typing import List, Tuple
 from math import ceil, log10
+from typing import (
+    List,
+    Tuple,
+    TYPE_CHECKING,
+)
 
 from poezio import poopt
 from poezio.theming import (
     get_theme,
 )
-from poezio.windows import Win
 from poezio.ui.consts import (
     FORMAT_CHAR,
     LONG_FORMAT,
@@ -27,6 +30,9 @@ from poezio.ui.types import (
     XMLLog,
 )
 
+if TYPE_CHECKING:
+    from poezio.windows import Win
+
 # msg is a reference to the corresponding Message object. text_start and
 # text_end are the position delimiting the text in this line.
 class Line:
@@ -110,7 +116,7 @@ def build_xmllog(msg: XMLLog, width: int, timestamp: bool, nick_size: int = 10)
 
 
 @singledispatch
-def write_pre(msg: BaseMessage, win: Win, with_timestamps: bool, nick_size: int) -> int:
+def write_pre(msg: BaseMessage, win: 'Win', with_timestamps: bool, nick_size: int) -> int:
     """Write the part before text (only the timestamp)"""
     if with_timestamps:
         return PreMessageHelpers.write_time(win, False, msg.time)
@@ -118,7 +124,7 @@ def write_pre(msg: BaseMessage, win: Win, with_timestamps: bool, nick_size: int)
 
 
 @write_pre.register(Message)
-def write_pre_message(msg: Message, win: Win, with_timestamps: bool, nick_size: int) -> int:
+def write_pre_message(msg: Message, win: 'Win', with_timestamps: bool, nick_size: int) -> int:
     """Write the part before the body:
         - timestamp (short or long)
         - ack/nack
@@ -162,7 +168,7 @@ def write_pre_message(msg: Message, win: Win, with_timestamps: bool, nick_size:
 
 
 @write_pre.register(XMLLog)
-def write_pre_xmllog(msg: XMLLog, win: Win, with_timestamps: bool, nick_size: int) -> int:
+def write_pre_xmllog(msg: XMLLog, win: 'Win', with_timestamps: bool, nick_size: int) -> int:
     """Write the part before the stanza (timestamp + IN/OUT)"""
     offset = 0
     if with_timestamps:
@@ -183,7 +189,7 @@ def write_pre_xmllog(msg: XMLLog, win: Win, with_timestamps: bool, nick_size: in
 class PreMessageHelpers:
 
     @staticmethod
-    def write_revisions(buffer: Win, msg: Message) -> int:
+    def write_revisions(buffer: 'Win', msg: Message) -> int:
         if msg.revisions:
             color = get_theme().COLOR_REVISIONS_MESSAGE
             with buffer.colored_text(color=color):
@@ -192,7 +198,7 @@ class PreMessageHelpers:
         return 0
 
     @staticmethod
-    def write_ack(buffer: Win) -> int:
+    def write_ack(buffer: 'Win') -> int:
         theme = get_theme()
         color = theme.COLOR_CHAR_ACK
         with buffer.colored_text(color=color):
@@ -201,7 +207,7 @@ class PreMessageHelpers:
         return poopt.wcswidth(theme.CHAR_ACK_RECEIVED) + 1
 
     @staticmethod
-    def write_nack(buffer: Win) -> int:
+    def write_nack(buffer: 'Win') -> int:
         theme = get_theme()
         color = theme.COLOR_CHAR_NACK
         with buffer.colored_text(color=color):
@@ -210,7 +216,7 @@ class PreMessageHelpers:
         return poopt.wcswidth(theme.CHAR_NACK) + 1
 
     @staticmethod
-    def write_nickname(buffer: Win, nickname: str, color, highlight=False) -> None:
+    def write_nickname(buffer: 'Win', nickname: str, color, highlight=False) -> None:
         """
         Write the nickname, using the user's color
         and return the number of written characters
@@ -228,7 +234,7 @@ class PreMessageHelpers:
             buffer.addstr(nickname)
 
     @staticmethod
-    def write_time(buffer: Win, history: bool, time: datetime) -> int:
+    def write_time(buffer: 'Win', history: bool, time: datetime) -> int:
         """
         Write the date on the yth line of the window
         """
diff --git a/poezio/windows/text_win.py b/poezio/windows/text_win.py
index 61ab4cc0..ce52e29d 100644
--- a/poezio/windows/text_win.py
+++ b/poezio/windows/text_win.py
@@ -318,11 +318,17 @@ class TextWin(Win):
         with_timestamps = config.get('show_timestamps')
         nick_size = config.get('max_nick_length')
         for i in range(len(self.built_lines) - 1, -1, -1):
-            if self.built_lines[i] and self.built_lines[i].msg.identifier == old_id:
+            current = self.built_lines[i]
+            if current is not None and current.msg.identifier == old_id:
                 index = i
-                while index >= 0 and self.built_lines[index] and self.built_lines[index].msg.identifier == old_id:
+                while (
+                        index >= 0
+                        and current is not None
+                        and current.msg.identifier == old_id
+                        ):
                     self.built_lines.pop(index)
                     index -= 1
+                current = self.built_lines[index]
                 index += 1
                 lines = build_lines(
                     message, self.width, timestamp=with_timestamps, nick_size=nick_size