Add badge to display unread count

Author: duynguyenhoang

app.py

@@ -5,13 +5,11 @@
 import json
 import os
 import requests
-import subprocess
 import sys
 import traceback
 import tempfile
 import urwid
 from datetime import datetime
-from slackclient import SlackClient
 from sclack.components import Attachment, Channel, ChannelHeader, ChatBox, Dm
 from sclack.components import Indicators, MarkdownText, Message, MessageBox
 from sclack.components import NewMessagesDivider, Profile, ProfileSideBar
@@ -36,6 +34,7 @@ def run(self):
     def set_exception_handler(self, handler):
         self._custom_exception_handler = handler
 
+
 class App:
     message_box = None
 
@@ -162,22 +161,51 @@ def mount_sidebar(self, executor):
         urwid.connect_signal(self.sidebar, 'go_to_channel', self.go_to_channel)
         loop.create_task(self.get_channels_info(executor, channels))
         loop.create_task(self.get_presences(executor, dms))
+        loop.create_task(self.get_dms_unread(executor, dms))
 
     @asyncio.coroutine
     def get_presences(self, executor, dm_widgets):
+        """
+        Compute and return presence because updating UI from another thread is unsafe
+        :param executor:
+        :param dm_widgets:
+        :return:
+        """
         def get_presence(dm_widget):
-            # Compute and return presence because updating UI from another thread is unsafe
             presence = self.store.get_presence(dm_widget.user)
             return [dm_widget, presence]
         presences = yield from asyncio.gather(*[
             loop.run_in_executor(executor, get_presence, dm_widget)
             for dm_widget in dm_widgets
         ])
+
         for presence in presences:
             [widget, response] = presence
             if response['ok']:
                 widget.set_presence(response['presence'])
 
+    @asyncio.coroutine
+    def get_dms_unread(self, executor, dm_widgets):
+        """
+        Compute and return unread_count_display because updating UI from another thread is unsafe
+        :param executor:
+        :param dm_widgets:
+        :return:
+        """
+        def get_presence(dm_widget):
+            profile_response = self.store.get_channel_info(dm_widget.id)
+            return [dm_widget, profile_response]
+
+        responses = yield from asyncio.gather(*[
+            loop.run_in_executor(executor, get_presence, dm_widget)
+            for dm_widget in dm_widgets
+        ])
+
+        for profile_response in responses:
+            [widget, response] = profile_response
+            if response is not None:
+                widget.set_unread(response['unread_count_display'])
+
     @asyncio.coroutine
     def get_channels_info(self, executor, channels):
         def get_info(channel):
@@ -187,10 +215,20 @@ def get_info(channel):
             loop.run_in_executor(executor, get_info, channel)
             for channel in channels
         ])
+
         for channel_info in channels_info:
             [widget, response] = channel_info
             widget.set_unread(response.get('unread_count_display', 0))
 
+    @asyncio.coroutine
+    def update_chat(self, event):
+        """
+        Update channel/DM message count badge
+        :param event:
+        :return:
+        """
+        self.sidebar.update_items(event)
+
     @asyncio.coroutine
     def mount_chatbox(self, executor, channel):
         yield from asyncio.gather(
@@ -406,7 +444,7 @@ def lazy_load_images(self, files, widget):
         if not self.config['features']['pictures']:
             return
 
-        allowed_file_types =  ('bmp', 'gif', 'jpeg', 'jpg', 'png')
+        allowed_file_types = ('bmp', 'gif', 'jpeg', 'jpg', 'png')
 
         for file in files:
             if file.get('filetype') in allowed_file_types:
@@ -417,7 +455,6 @@ def lazy_load_images(self, files, widget):
                     not file.get('is_external', True)
                 ))
 
-
     def render_messages(self, messages):
         _messages = []
         previous_date = self.store.state.last_date
@@ -435,6 +472,7 @@ def render_messages(self, messages):
                     date_text = 'Today'
                 else:
                     date_text = message_date.strftime('%A, %B %d')
+
             # New messages badge
             if (message_datetime > last_read_datetime and not self.store.state.did_render_new_messages
                 and (self.store.state.channel.get('unread_count_display', 0) > 0)):
@@ -527,7 +565,10 @@ def start_real_time(self):
         self.store.slack.rtm_connect(auto_reconnect=True)
 
         def stop_typing(*args):
-            self.chatbox.message_box.typing = None
+            # Prevent error while switching workspace
+            if self.chatbox is not None:
+                self.chatbox.message_box.typing = None
+
         alarm = None
         while self.store.slack.server.connected is True:
             events = self.store.slack.rtm_read()
@@ -539,9 +580,12 @@ def stop_typing(*args):
                     for channel in self.sidebar.channels:
                         if channel.id == event['channel']:
                             channel.set_unread(unread)
-                elif event.get('channel') == self.store.state.channel['id']:
-                    if event['type'] == 'message':
-                        # Delete message
+                elif event['type'] == 'message':
+                    loop.create_task(
+                        self.update_chat(event)
+                    )
+
+                    if event.get('channel') == self.store.state.channel['id']:
                         if event.get('subtype') == 'message_deleted':
                             for widget in self.chatbox.body.body:
                                 if hasattr(widget, 'ts') and getattr(widget, 'ts') == event['deleted_ts']: