Fetch multiple users/bots in the same request

Turns out you can query for multiple users/bots with
users.info/bots.info even though it's not documented. Use that to query
for multiple users/bots in one request, instead of making tons of
requests.

5 files changed, 145 insertions, 40 deletions

diff --git a/slack/slack_api.py b/slack/slack_api.py
index b13d27c..656227e 100644
--- a/slack/slack_api.py
+++ b/slack/slack_api.py
@@ -1,22 +1,21 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING, Dict, Union
+from typing import TYPE_CHECKING, Dict, List, Union
 from urllib.parse import urlencode
 
 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_bots_info import SlackBotInfoResponse, SlackBotsInfoResponse
     from slack_api.slack_conversations_history import SlackConversationsHistoryResponse
     from slack_api.slack_conversations_info import SlackConversationsInfoResponse
     from slack_api.slack_rtm_connect import SlackRtmConnectResponse
     from slack_api.slack_users_conversations import SlackUsersConversationsResponse
-    from slack_api.slack_users_info import SlackUserInfoResponse
+    from slack_api.slack_users_info import SlackUserInfoResponse, SlackUsersInfoResponse
 
     from slack.slack_conversation import SlackConversation
-    from slack.slack_user import SlackBot, SlackUser
     from slack.slack_workspace import SlackWorkspace
 
 
@@ -87,8 +86,14 @@ class SlackApi:
             pages,
         )
 
-    async def fetch_user_info(self, user: SlackUser) -> SlackUserInfoResponse:
-        return await self._fetch("users.info", {"user": user.id})
+    async def fetch_user_info(self, user_id: str) -> SlackUserInfoResponse:
+        return await self._fetch("users.info", {"user": user_id})
 
-    async def fetch_bot_info(self, bot: SlackBot) -> SlackBotInfoResponse:
-        return await self._fetch("bots.info", {"bot": bot.id})
+    async def fetch_users_info(self, user_ids: List[str]) -> SlackUsersInfoResponse:
+        return await self._fetch("users.info", {"users": ",".join(user_ids)})
+
+    async def fetch_bot_info(self, bot_id: str) -> SlackBotInfoResponse:
+        return await self._fetch("bots.info", {"bot": bot_id})
+
+    async def fetch_bots_info(self, bot_ids: List[str]) -> SlackBotsInfoResponse:
+        return await self._fetch("bots.info", {"bots": ",".join(bot_ids)})
diff --git a/slack/slack_conversation.py b/slack/slack_conversation.py
index d1e5707..8123df8 100644
--- a/slack/slack_conversation.py
+++ b/slack/slack_conversation.py
@@ -90,6 +90,10 @@ class SlackConversation:
             start = time.time()
 
             messages = [SlackMessage(self, message) for message in history["messages"]]
+
+            sender_user_ids = [m.sender_user_id for m in messages if m.sender_user_id]
+            await self.workspace.users.initialize_items(sender_user_ids)
+
             messages_rendered = await gather(
                 *(message.render_message() for message in messages)
             )
diff --git a/slack/slack_message.py b/slack/slack_message.py
index 3c5383b..0757761 100644
--- a/slack/slack_message.py
+++ b/slack/slack_message.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import re
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, List, Optional
 
 from slack.slack_user import format_bot_nick
 from slack.task import gather
@@ -23,26 +23,32 @@ class SlackMessage:
     def workspace(self) -> SlackWorkspace:
         return self.conversation.workspace
 
-    async def render_message(self):
+    @property
+    def sender_user_id(self) -> Optional[str]:
+        return self._message_json.get("user")
+
+    async def render_message(self) -> str:
+        prefix_coro = self._prefix()
+        message_coro = self._unfurl_refs(self._message_json["text"])
+        prefix, message = await gather(prefix_coro, message_coro)
+        return f"{prefix}\t{message}"
+
+    async def _prefix(self) -> str:
         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)
+                return format_bot_nick(username, colorize=True)
             else:
                 bot = await self.workspace.bots[self._message_json["bot_id"]]
-                prefix = bot.nick(colorize=True)
+                return bot.nick(colorize=True)
         else:
             user = await self.workspace.users[self._message_json["user"]]
-            prefix = user.nick(colorize=True)
-
-        message = await self._unfurl_refs(self._message_json["text"])
-
-        return f"{prefix}\t{message}"
+            return user.nick(colorize=True)
 
-    async def _unfurl_refs(self, message: str):
+    async def _unfurl_refs(self, message: str) -> str:
         re_user = re.compile("<@([^>]+)>")
         user_ids: List[str] = re_user.findall(message)
         users_list = await gather(
diff --git a/slack/slack_user.py b/slack/slack_user.py
index f5b6d1f..6b9b056 100644
--- a/slack/slack_user.py
+++ b/slack/slack_user.py
@@ -1,13 +1,17 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 import weechat
 
 from slack.shared import shared
+from slack.task import create_task
 from slack.util import with_color
 
 if TYPE_CHECKING:
+    from slack_api.slack_bots_info import SlackBotInfo
+    from slack_api.slack_users_info import SlackUserInfo
+
     from slack.slack_workspace import SlackApi, SlackWorkspace
 
 
@@ -25,20 +29,34 @@ def format_bot_nick(nick: str, colorize: bool = False) -> str:
 
 
 class SlackUser:
-    def __init__(self, workspace: SlackWorkspace, id: str):
+    def __init__(
+        self,
+        workspace: SlackWorkspace,
+        id: str,
+        info: Optional[SlackUserInfo] = None,
+    ):
         self.workspace = workspace
         self.id = id
+        if info:
+            self._info = info
+            self._set_info_task = None
+        else:
+            self._set_info_task = create_task(self._set_info())
 
     @property
     def _api(self) -> SlackApi:
         return self.workspace.api
 
-    async def init(self):
-        info = await self._api.fetch_user_info(self)
-        if info["ok"] is False:
+    async def _set_info(self):
+        info_response = await self._api.fetch_user_info(self.id)
+        if info_response["ok"] is False:
             # TODO: Handle error
             raise Exception("Failed fetching user info")
-        self._info = info["user"]
+        self._info = info_response["user"]
+
+    async def ensure_initialized(self):
+        if self._set_info_task:
+            await self._set_info_task
 
     def nick(self, colorize: bool = False) -> str:
         nick = self._name_without_spaces()
@@ -75,20 +93,34 @@ class SlackUser:
 
 
 class SlackBot:
-    def __init__(self, workspace: SlackWorkspace, id: str):
+    def __init__(
+        self,
+        workspace: SlackWorkspace,
+        id: str,
+        info: Optional[SlackBotInfo] = None,
+    ):
         self.workspace = workspace
         self.id = id
+        if info:
+            self._info = info
+            self._set_info_task = None
+        else:
+            self._set_info_task = create_task(self._set_info())
 
     @property
     def _api(self) -> SlackApi:
         return self.workspace.api
 
-    async def init(self):
-        info = await self._api.fetch_bot_info(self)
-        if info["ok"] is False:
+    async def _set_info(self):
+        info_response = await self._api.fetch_bot_info(self.id)
+        if info_response["ok"] is False:
             # TODO: Handle error
-            raise Exception("Failed fetching user info")
-        self._info = info["bot"]
+            raise Exception("Failed fetching bot info")
+        self._info = info_response["bot"]
+
+    async def ensure_initialized(self):
+        if self._set_info_task:
+            await self._set_info_task
 
     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 77c9e46..bf0f13f 100644
--- a/slack/slack_workspace.py
+++ b/slack/slack_workspace.py
@@ -4,7 +4,7 @@ import json
 import socket
 import ssl
 import time
-from typing import Any, Dict, Generic, Type, TypeVar
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import weechat
 from websocket import ABNF, WebSocketConnectionClosedException, create_connection
@@ -17,24 +17,82 @@ from slack.slack_user import SlackBot, SlackUser
 from slack.task import Future, create_task
 from slack.util import get_callback_name
 
-SlackUserOrBot = TypeVar("SlackUserOrBot", SlackUser, SlackBot)
+if TYPE_CHECKING:
+    from slack_api.slack_bots_info import SlackBotInfo
+    from slack_api.slack_users_info import SlackUserInfo
 
 
-class SlackUsersOrBots(Generic[SlackUserOrBot], Dict[str, Future[SlackUserOrBot]]):
-    def __init__(self, workspace: SlackWorkspace, item_class: Type[SlackUserOrBot]):
+class SlackUsers(Dict[str, Future[SlackUser]]):
+    def __init__(self, workspace: SlackWorkspace):
         super().__init__()
         self.workspace = workspace
-        self.item_class = item_class
 
     def __missing__(self, key: str):
         self[key] = create_task(self._create_item(key))
         return self[key]
 
-    async def _create_item(self, item_id: str) -> SlackUserOrBot:
-        item = self.item_class(self.workspace, item_id)
-        await item.init()
+    async def initialize_items(self, item_ids: List[str]):
+        items_info_task = create_task(self._fetch_items_info(item_ids))
+        for item_id in set(item_ids):
+            self[item_id] = create_task(self._create_item(item_id, items_info_task))
+
+    async def _create_item(
+        self,
+        item_id: str,
+        items_info_task: Optional[Future[Dict[str, SlackUserInfo]]] = None,
+    ) -> SlackUser:
+        if items_info_task:
+            items_info = await items_info_task
+            item = SlackUser(self.workspace, item_id, items_info[item_id])
+        else:
+            item = SlackUser(self.workspace, item_id)
+
+        await item.ensure_initialized()
         return item
 
+    async def _fetch_items_info(self, item_ids: List[str]):
+        response = await self.workspace.api.fetch_users_info(item_ids)
+        if response["ok"] is False:
+            # TODO: Handle error
+            raise Exception("Failed fetching users")
+        return {info["id"]: info for info in response["users"]}
+
+
+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_item(key))
+        return self[key]
+
+    async def initialize_items(self, item_ids: List[str]):
+        items_info_task = create_task(self._fetch_items_info(item_ids))
+        for item_id in set(item_ids):
+            self[item_id] = create_task(self._create_item(item_id, items_info_task))
+
+    async def _create_item(
+        self,
+        item_id: str,
+        items_info_task: Optional[Future[Dict[str, SlackBotInfo]]] = None,
+    ) -> SlackBot:
+        if items_info_task:
+            items_info = await items_info_task
+            item = SlackBot(self.workspace, item_id, items_info[item_id])
+        else:
+            item = SlackBot(self.workspace, item_id)
+
+        await item.ensure_initialized()
+        return item
+
+    async def _fetch_items_info(self, item_ids: List[str]):
+        response = await self.workspace.api.fetch_bots_info(item_ids)
+        if response["ok"] is False:
+            # TODO: Handle error
+            raise Exception("Failed fetching bots")
+        return {info["id"]: info for info in response["bots"]}
+
 
 class SlackWorkspace:
     def __init__(self, name: str):
@@ -42,8 +100,8 @@ class SlackWorkspace:
         self.config = shared.config.create_workspace_config(self.name)
         self.api = SlackApi(self)
         self._is_connected = False
-        self.users = SlackUsersOrBots(self, SlackUser)
-        self.bots = SlackUsersOrBots(self, SlackBot)
+        self.users = SlackUsers(self)
+        self.bots = SlackBots(self)
         self.conversations: Dict[str, SlackConversation] = {}
 
     @property