From 96d23510cde5fe96b334961a92c5bc3944a6973a Mon Sep 17 00:00:00 2001
From: Alexander Zveruk <sanyazveruk@gmail.com>
Date: Tue, 19 May 2020 23:29:50 +0300
Subject: [PATCH] fix bug with message update wrong type in model

---
 tg/controllers/__init__.py | 33 ++++++++++++++-------------------
 tg/msg.py                  |  8 ++++++--
 2 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/tg/controllers/__init__.py b/tg/controllers/__init__.py
index 52375ce..b2aa726 100644
--- a/tg/controllers/__init__.py
+++ b/tg/controllers/__init__.py
@@ -321,29 +321,29 @@ class Controller:
         self.view.msgs.draw(current_msg_idx, msgs, MSGS_LEFT_SCROLL_THRESHOLD)
 
     @handle_exception
-    def update_msg_content(self, update):
-        content = MsgProxy(update["new_content"])
+    def update_msg_content(self, update: Dict[str, Any]):
         chat_id = update["chat_id"]
         message_id = update["message_id"]
-        self.model.msgs.update_msg_content(chat_id, message_id, content)
+        self.model.msgs.update_msg_content(
+            chat_id, message_id, update["new_content"]
+        )
         current_chat_id = self.model.chats.id_by_index(self.model.current_chat)
         if current_chat_id == chat_id:
             self.refresh_msgs()
 
     @handle_exception
-    def update_new_msg(self, update):
+    def update_new_msg(self, update: Dict[str, Any]):
         msg = MsgProxy(update["message"])
-        chat_id = msg["chat_id"]
-        self.model.msgs.add_message(chat_id, msg)
+        self.model.msgs.add_message(msg.chat_id, msg.msg)
         current_chat_id = self.model.chats.id_by_index(self.model.current_chat)
-        if current_chat_id == chat_id:
+        if current_chat_id == msg.chat_id:
             self.refresh_msgs()
         if msg.file_id and msg.size <= config.max_download_size:
-            self.download(msg.file_id, chat_id, msg["id"])
+            self.download(msg.file_id, msg.chat_id, msg["id"])
 
-        self._notify_for_message(chat_id, msg)
+        self._notify_for_message(msg.chat_id, msg)
 
-    def _notify_for_message(self, chat_id: int, msg: Dict[str, Any]):
+    def _notify_for_message(self, chat_id: int, msg: MsgProxy):
         # do not notify, if muted
         # TODO: optimize
         chat = None
@@ -356,17 +356,12 @@ class Controller:
             return
 
         # notify
-        user_id = msg["sender_user_id"]
-        if msg["sender_user_id"] == self.model.get_me()["id"]:
+        if msg.sender_id == self.model.get_me()["id"]:
             return
-        user = self.model.users.get_user(user_id)
-        name = "{} {}".format(user["first_name"], user["last_name"])
-        _type = msg["content"]["@type"]
+        user = self.model.users.get_user(msg.sender_id)
+        name = f"{user['first_name']} {user['last_name']}"
 
-        if _type == "messageText":
-            text = msg["content"]["text"]["text"]
-        else:
-            text = MsgProxy.types.get(_type, "")
+        text = msg.text_content if msg.is_text else msg.content_type
         notify(text, title=name)
 
     @handle_exception
diff --git a/tg/msg.py b/tg/msg.py
index 76e9ba2..363909a 100644
--- a/tg/msg.py
+++ b/tg/msg.py
@@ -135,6 +135,10 @@ class MsgProxy:
     def is_text(self):
         return self.msg["content"]["@type"] == "messageText"
 
+    @property
+    def text_content(self) -> str:
+        return self.msg["content"]["text"]["text"]
+
     @property
     def is_downloaded(self):
         doc = self.get_doc(self.msg)
@@ -145,8 +149,8 @@ class MsgProxy:
         return self.msg.get("reply_to_message_id")
 
     @property
-    def chat_id(self):
-        return self.msg.get("chat_id")
+    def chat_id(self) -> int:
+        return self.msg["chat_id"]
 
     @property
     def sender_id(self) -> int: