Combine SlackUsers and SlackBots

1 files changed, 12 insertions, 24 deletions

diff --git a/slack/slack_workspace.py b/slack/slack_workspace.py
index ba79281..f2eedb0 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
+from typing import Any, Dict, Generic, Type, TypeVar
 
 import weechat
 from websocket import ABNF, WebSocketConnectionClosedException, create_connection
@@ -17,35 +17,23 @@ from slack.slack_user import SlackBot, SlackUser
 from slack.task import Future, create_task
 from slack.util import get_callback_name
 
-
-class SlackUsers(Dict[str, Future[SlackUser]]):
-    def __init__(self, workspace: SlackWorkspace):
-        super().__init__()
-        self.workspace = workspace
-
-    def __missing__(self, key: str):
-        self[key] = create_task(self._create_user(key))
-        return self[key]
-
-    async def _create_user(self, user_id: str) -> SlackUser:
-        user = SlackUser(self.workspace, user_id)
-        await user.init()
-        return user
+SlackUserOrBot = TypeVar("SlackUserOrBot", SlackUser, SlackBot)
 
 
-class SlackBots(Dict[str, Future[SlackBot]]):
-    def __init__(self, workspace: SlackWorkspace):
+class SlackUsersOrBots(Generic[SlackUserOrBot], Dict[str, Future[SlackUserOrBot]]):
+    def __init__(self, workspace: SlackWorkspace, item_class: Type[SlackUserOrBot]):
         super().__init__()
         self.workspace = workspace
+        self.item_class = item_class
 
     def __missing__(self, key: str):
-        self[key] = create_task(self._create_bot(key))
+        self[key] = create_task(self._create_item(key))
         return self[key]
 
-    async def _create_bot(self, bot_id: str) -> SlackBot:
-        bot = SlackBot(self.workspace, bot_id)
-        await bot.init()
-        return bot
+    async def _create_item(self, item_id: str) -> SlackUserOrBot:
+        item = self.item_class(self.workspace, item_id)
+        await item.init()
+        return item
 
 
 class SlackWorkspace:
@@ -54,8 +42,8 @@ class SlackWorkspace:
         self.config = shared.config.create_workspace_config(self.name)
         self.api = SlackApi(self)
         self.is_connected = False
-        self.users = SlackUsers(self)
-        self.bots = SlackBots(self)
+        self.users = SlackUsersOrBots(self, SlackUser)
+        self.bots = SlackUsersOrBots(self, SlackBot)
         self.conversations: Dict[str, SlackConversation] = {}
 
     def __setattr__(self, __name: str, __value: Any) -> None: