refactor: helper: Extract updating unread to function in set_count.

Having a nested function allows us to discontinue the use of the `key`
variable which can be useful if multiple sections of unread_counts need
to be updated for a message type.

It includes type annotation changes such as:

* `key` passed into the nested function is typed using TypeVar.
* Message has been typed using a TypedDict so that mypy can determine
  the type of `key`. Note the use of Totality here because stream
  messages/pms/huddles do not have the same attributes. This change
  cannot be split into a prior commit because that will require
  multiple type definitions of the same `key` variable.
* An alternate declaration of `messages` in initial_index is necessary
  due to a mypy bug - python/mypy#7217 - while
  using TypeDict with defaultdict.

Author: kaustubh-nair

zulipterminal/helper.py

@@ -21,7 +21,29 @@
 LINUX = platform.system() == "Linux"
 WSL = 'microsoft' in platform.release().lower()
 
-Message = Dict[str, Any]
+Message = TypedDict('Message', {
+    'id': int,
+    'sender_id': int,
+    'content': str,
+    'recipient_id': int,
+    'timestamp': int,
+    'client': str,
+    'subject': str,
+    'topic_links': List[str],
+    'is_me_message': bool,
+    'reactions': List[Dict[str, Any]],
+    'submessages': List[Dict[str, Any]],
+    'flags': List[str],
+    'sender_full_name': str,
+    'sender_short_name': str,
+    'sender_email': str,
+    'sender_realm_str': str,
+    'display_recipient': Any,
+    'type': str,
+    'stream_id': int,
+    'avatar_url': str,
+    'content_type': str,
+  }, total=False)
 
 Index = TypedDict('Index', {
     'pointer': Dict[str, Union[int, Set[None]]],  # narrow_str, message_id
@@ -53,7 +75,7 @@
     edited_messages=set(),
     topics=defaultdict(list),
     search=set(),
-    messages=defaultdict(dict),
+    messages=defaultdict(lambda: Message()),
 )
 
 
@@ -92,30 +114,32 @@ def _set_count_in_model(new_count: int, changed_messages: List[Message],
         but updates `unread_counts` (which can update the model
         if it's passed in, but is not tied to it).
     """
-    for message in changed_messages:
-        if message['type'] == 'stream':
-            key = (message['stream_id'], message['subject'])
-            unreads = unread_counts['unread_topics']
-        # self-pm has only one display_recipient
-        # 1-1 pms have 2 display_recipient
-        elif len(message['display_recipient']) <= 2:
-            key = message['sender_id']
-            unreads = unread_counts['unread_pms']  # type: ignore
-        else:  # If it's a group pm
-            key = frozenset(  # type: ignore
-                recipient['id'] for recipient in message['display_recipient']
-            )
-            unreads = unread_counts['unread_huddles']  # type: ignore
+    # broader unread counts (for all_*) are updated
+    # later conditionally in _set_count_in_view.
+    KeyT = TypeVar('KeyT')
 
-        # broader unread counts (for all_* and streams) are updated
-        # later conditionally.
+    def update_unreads(unreads: Dict[KeyT, int], key: KeyT) -> None:
         if key in unreads:
             unreads[key] += new_count
             if unreads[key] == 0:
                 unreads.pop(key)
         elif new_count == 1:
             unreads[key] = new_count
 
+    for message in changed_messages:
+        if message['type'] == 'stream':
+            stream_id = message['stream_id']
+            update_unreads(unread_counts['unread_topics'],
+                           (stream_id, message['subject']))
+        # self-pm has only one display_recipient
+        # 1-1 pms have 2 display_recipient
+        elif len(message['display_recipient']) <= 2:
+            update_unreads(unread_counts['unread_pms'], message['sender_id'])
+        else:  # If it's a group pm
+            update_unreads(unread_counts['unread_huddles'],
+                           frozenset(recipient['id'] for recipient
+                                     in message['display_recipient']))
+
 
 def _set_count_in_view(controller: Any, new_count: int,
                        changed_messages: List[Message],