Merge function with classify_unread_counts

Author: kaustubh-nair

zulipterminal/helper.py

@@ -38,7 +38,7 @@
     # Downloaded message data
     'messages': Dict[int, Message],  # message_id: Message
     # unread message data; additional data in model.initial_data['unread_msgs']
-    'unread_msgs': Set[int]  # {message_ids, ...}
+    'unread_msg_ids': Set[int]  # {message_ids, ...}
 })
 
 initial_index = Index(
@@ -53,7 +53,7 @@
     topics=defaultdict(list),
     search=set(),
     messages=defaultdict(dict),
-    unread_msgs=set(),
+    unread_msg_ids=set(),
 )
 
 
@@ -353,26 +353,10 @@ def index_messages(messages: List[Any],
     return index
 
 
-def index_unread_messsages(unread_msgs: Dict[str, Any], index: Index) -> Index:
-
-    for pm in unread_msgs['pms']:
-        message_ids = pm['unread_message_ids']
-        index['unread_msgs'].update(message_ids)
-
-    for stream in unread_msgs['streams']:
-        message_ids = stream['unread_message_ids']
-        index['unread_msgs'].update(message_ids)
-
-    for group_pm in unread_msgs['huddles']:
-        message_ids = group_pm['unread_message_ids']
-        index['unread_msgs'].update(message_ids)
-
-    return index
-
-
 def classify_unread_counts(model: Any) -> UnreadCounts:
     # TODO: support group pms
     unread_msg_counts = model.initial_data['unread_msgs']
+    index = model.index
 
     unread_counts = UnreadCounts(
         all_msg=0,
@@ -383,14 +367,19 @@ def classify_unread_counts(model: Any) -> UnreadCounts:
         streams=dict(),
     )
 
+    # add message_ids to index and count to UnreadCounts
     for pm in unread_msg_counts['pms']:
-        count = len(pm['unread_message_ids'])
+        message_ids = pm['unread_message_ids']
+        index['unread_msg_ids'].update(message_ids)
+        count = len(message_ids)
         unread_counts['unread_pms'][pm['sender_id']] = count
         unread_counts['all_msg'] += count
         unread_counts['all_pms'] += count
 
     for stream in unread_msg_counts['streams']:
-        count = len(stream['unread_message_ids'])
+        message_ids = stream['unread_message_ids']
+        index['unread_msg_ids'].update(message_ids)
+        count = len(message_ids)
         stream_id = stream['stream_id']
         if [model.stream_dict[stream_id]['name'],
                 stream['topic']] in model.muted_topics:
@@ -405,7 +394,9 @@ def classify_unread_counts(model: Any) -> UnreadCounts:
 
     # store unread count of group pms in `unread_huddles`
     for group_pm in unread_msg_counts['huddles']:
-        count = len(group_pm['unread_message_ids'])
+        message_ids = group_pm['unread_message_ids']
+        index['unread_msg_ids'].update(message_ids)
+        count = len(message_ids)
         user_ids = group_pm['user_ids_string'].split(',')
         user_ids = frozenset(map(int, user_ids))
         unread_counts['unread_huddles'][user_ids] = count

zulipterminal/model.py

@@ -22,7 +22,6 @@
     initial_index,
     Message,
     notify,
-    index_unread_messsages,
 )
 from zulipterminal.ui_tools.utils import create_msg_box_list
 
@@ -287,7 +286,7 @@ def toggle_message_star_status(self, message: Dict[str, Any]) -> None:
     def mark_message_ids_as_read(self, id_list: List[int]) -> None:
         if not id_list:
             return
-        self.index['unread_msgs'].difference_update(id_list)
+        self.index['unread_msg_ids'].difference_update(id_list)
         self.client.update_message_flags({
             'messages': id_list,
             'flag': 'read',
@@ -448,9 +447,6 @@ def _update_initial_data(self) -> None:
             self.user_email = self.initial_data['email']
             self.user_full_name = self.initial_data['full_name']
             self.server_name = self.initial_data['realm_name']
-            unread_msgs = self.initial_data['unread_msgs']
-            self.index = index_unread_messsages(unread_msgs,
-                                                self.index)
         else:
             failures = [name for name, result in results.items() if not result]
             raise ServerConnectionFailure(", ".join(failures))