Show nick suffix for bots

7 files changed, 121 insertions, 11 deletions

diff --git a/slack/config.py b/slack/config.py
index 56b828f..4533780 100644
--- a/slack/config.py
+++ b/slack/config.py
@@ -48,6 +48,13 @@ class SlackConfigSectionLook:
     def __init__(self, weechat_config: WeeChatConfig):
         self._section = WeeChatSection(weechat_config, "look")
 
+        self.bot_user_suffix = WeeChatOption(
+            self._section,
+            "bot_user_suffix",
+            "the suffix appended to nicks to indicate a bot",
+            " :]",
+        )
+
         self.external_user_suffix = WeeChatOption(
             self._section,
             "external_user_suffix",
diff --git a/slack/slack_api.py b/slack/slack_api.py
index 463dce9..081a2f4 100644
--- a/slack/slack_api.py
+++ b/slack/slack_api.py
@@ -8,6 +8,7 @@ from slack.http import http_request
 from slack.shared import shared
 
 if TYPE_CHECKING:
+    from slack_api.slack_bots_info import SlackBotInfoResponse
     from slack_api.slack_conversations_history import SlackConversationsHistoryResponse
     from slack_api.slack_conversations_info import SlackConversationsInfoResponse
     from slack_api.slack_rtm_connect import SlackRtmConnectResponse
@@ -15,7 +16,7 @@ if TYPE_CHECKING:
     from slack_api.slack_users_info import SlackUsersInfoResponse
 
     from slack.slack_conversation import SlackConversation
-    from slack.slack_user import SlackUser
+    from slack.slack_user import SlackBot, SlackUser
     from slack.slack_workspace import SlackWorkspace
 
 
@@ -88,3 +89,6 @@ class SlackApi:
 
     async def fetch_users_info(self, user: SlackUser) -> SlackUsersInfoResponse:
         return await self._fetch("users.info", {"user": user.id})
+
+    async def fetch_bots_info(self, bot: SlackBot) -> SlackBotInfoResponse:
+        return await self._fetch("bots.info", {"bot": bot.id})
diff --git a/slack/slack_message.py b/slack/slack_message.py
index d6431d7..3c5383b 100644
--- a/slack/slack_message.py
+++ b/slack/slack_message.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import re
 from typing import TYPE_CHECKING, List
 
+from slack.slack_user import format_bot_nick
 from slack.task import gather
 
 if TYPE_CHECKING:
@@ -23,12 +24,21 @@ class SlackMessage:
         return self.conversation.workspace
 
     async def render_message(self):
-        message = await self._unfurl_refs(self._message_json["text"])
-        if "user" in self._message_json:
+        if (
+            "subtype" in self._message_json
+            and self._message_json["subtype"] == "bot_message"
+        ):
+            username = self._message_json.get("username")
+            if username:
+                prefix = format_bot_nick(username, colorize=True)
+            else:
+                bot = await self.workspace.bots[self._message_json["bot_id"]]
+                prefix = bot.nick(colorize=True)
+        else:
             user = await self.workspace.users[self._message_json["user"]]
             prefix = user.nick(colorize=True)
-        else:
-            prefix = "bot"
+
+        message = await self._unfurl_refs(self._message_json["text"])
 
         return f"{prefix}\t{message}"
 
diff --git a/slack/slack_user.py b/slack/slack_user.py
index 7b04605..4d2c29c 100644
--- a/slack/slack_user.py
+++ b/slack/slack_user.py
@@ -11,6 +11,19 @@ if TYPE_CHECKING:
     from slack.slack_workspace import SlackApi, SlackWorkspace
 
 
+def nick_color(nick: str) -> str:
+    return weechat.info_get("nick_color_name", nick)
+
+
+def format_bot_nick(nick: str, colorize: bool = False) -> str:
+    nick = nick.replace(" ", "")
+
+    if colorize:
+        nick = with_color(nick_color(nick), nick)
+
+    return nick + shared.config.look.bot_user_suffix.value
+
+
 class SlackUser:
     def __init__(self, workspace: SlackWorkspace, id: str):
         self.workspace = workspace
@@ -58,4 +71,24 @@ class SlackUser:
                 weechat.config_get("weechat.color.chat_nick_self")
             )
 
-        return weechat.info_get("nick_color_name", self._name_without_spaces())
+        return nick_color(self._name_without_spaces())
+
+
+class SlackBot:
+    def __init__(self, workspace: SlackWorkspace, id: str):
+        self.workspace = workspace
+        self.id = id
+
+    @property
+    def _api(self) -> SlackApi:
+        return self.workspace.api
+
+    async def init(self):
+        info = await self._api.fetch_bots_info(self)
+        if info["ok"] is False:
+            # TODO: Handle error
+            raise Exception("Failed fetching user info")
+        self._info = info["bot"]
+
+    def nick(self, colorize: bool = False) -> str:
+        return format_bot_nick(self._info["name"], colorize)
diff --git a/slack/slack_workspace.py b/slack/slack_workspace.py
index 4e18634..ba79281 100644
--- a/slack/slack_workspace.py
+++ b/slack/slack_workspace.py
@@ -13,7 +13,7 @@ from slack.proxy import Proxy
 from slack.shared import shared
 from slack.slack_api import SlackApi
 from slack.slack_conversation import SlackConversation
-from slack.slack_user import SlackUser
+from slack.slack_user import SlackBot, SlackUser
 from slack.task import Future, create_task
 from slack.util import get_callback_name
 
@@ -33,6 +33,21 @@ class SlackUsers(Dict[str, Future[SlackUser]]):
         return user
 
 
+class SlackBots(Dict[str, Future[SlackBot]]):
+    def __init__(self, workspace: SlackWorkspace):
+        super().__init__()
+        self.workspace = workspace
+
+    def __missing__(self, key: str):
+        self[key] = create_task(self._create_bot(key))
+        return self[key]
+
+    async def _create_bot(self, bot_id: str) -> SlackBot:
+        bot = SlackBot(self.workspace, bot_id)
+        await bot.init()
+        return bot
+
+
 class SlackWorkspace:
     def __init__(self, name: str):
         self.name = name
@@ -40,6 +55,7 @@ class SlackWorkspace:
         self.api = SlackApi(self)
         self.is_connected = False
         self.users = SlackUsers(self)
+        self.bots = SlackBots(self)
         self.conversations: Dict[str, SlackConversation] = {}
 
     def __setattr__(self, __name: str, __value: Any) -> None:
diff --git a/typings/slack_api/slack_bots_info.pyi b/typings/slack_api/slack_bots_info.pyi
new file mode 100644
index 0000000..0fb296f
--- /dev/null
+++ b/typings/slack_api/slack_bots_info.pyi
@@ -0,0 +1,24 @@
+from typing import Dict, Literal, TypedDict, final
+
+from typing_extensions import NotRequired
+
+class SlackBotInfo(TypedDict):
+    id: str
+    deleted: bool
+    name: str
+    updated: int
+    app_id: str
+    user_id: NotRequired[str]
+    icons: Dict[str, str]
+
+@final
+class SlackBotsInfoErrorResponse(TypedDict):
+    ok: Literal[False]
+    error: str
+
+@final
+class SlackBotsInfoSuccessResponse(TypedDict):
+    ok: Literal[True]
+    bot: SlackBotInfo
+
+SlackBotInfoResponse = SlackBotsInfoSuccessResponse | SlackBotsInfoErrorResponse
diff --git a/typings/slack_api/slack_conversations_history.pyi b/typings/slack_api/slack_conversations_history.pyi
index 479fe50..b0220e3 100644
--- a/typings/slack_api/slack_conversations_history.pyi
+++ b/typings/slack_api/slack_conversations_history.pyi
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import List, Literal, TypedDict, final
+from typing import Dict, List, Literal, TypedDict, final
 
 from typing_extensions import NotRequired
 
@@ -81,17 +81,21 @@ class SlackMessageFile(TypedDict):
 class SlackMessageCommon(TypedDict):
     type: Literal["message"]
     text: str
-    user: str
     ts: str
     reactions: NotRequired[List[SlackMessageReaction]]
 
-class SlackMessageStandard(SlackMessageCommon):
+class SlackMessageStandardCommon(SlackMessageCommon):
     client_msg_id: NotRequired[str]
+    user: str
     blocks: List[SlackMessageBlock]
     attachments: NotRequired[List[SlackMessageAttachment]]
     team: str
 
-class SlackMessageThreadParentCommon(SlackMessageStandard):
+@final
+class SlackMessageStandard(SlackMessageStandardCommon):
+    pass
+
+class SlackMessageThreadParentCommon(SlackMessageStandardCommon):
     thread_ts: str
     reply_count: int
     reply_users_count: int
@@ -108,21 +112,32 @@ class SlackMessageThreadParentSubscribed(SlackMessageThreadParentCommon):
     subscribed: Literal[True]
     last_read: str
 
+@final
 class SlackMessageWithFiles(SlackMessageCommon):
+    user: str
     files: List[SlackMessageFile]
     upload: bool
     display_as_bot: bool
 
 # TODO: Add other subtypes
 @final
+class SlackMessageSubtypeBotMessage(SlackMessageCommon):
+    subtype: Literal["bot_message"]
+    bot_id: str
+    username: NotRequired[str]
+    icons: NotRequired[Dict[str, str]]
+
+@final
 class SlackMessageSubtypeBotRemove(SlackMessageCommon):
     subtype: Literal["bot_remove"]
+    user: str
     bot_id: str
     bot_link: str
 
 @final
 class SlackMessageSubtypeBotAdd(SlackMessageCommon):
     subtype: Literal["bot_add"]
+    user: str
     bot_id: str
     bot_link: str
 
@@ -136,6 +151,7 @@ SlackMessage = (
     | SlackMessageThreadParentNotSubscribed
     | SlackMessageThreadParentSubscribed
     | SlackMessageWithFiles
+    | SlackMessageSubtypeBotMessage
     | SlackMessageSubtypeBotRemove
     | SlackMessageSubtypeBotAdd
 )