Get user from dict instead of by function

2 files changed, 18 insertions, 15 deletions

diff --git a/slack/slack_message.py b/slack/slack_message.py
index d90ddb8..0a7d9d2 100644
--- a/slack/slack_message.py
+++ b/slack/slack_message.py
@@ -27,7 +27,7 @@ class SlackMessage:
     async def render_message(self):
         message = await self.unfurl_refs(self.message_json["text"])
         if "user" in self.message_json:
-            user = await self.workspace.get_user(self.message_json["user"])
+            user = await self.workspace.users[self.message_json["user"]]
             prefix = user.name
         else:
             prefix = "bot"
@@ -38,7 +38,7 @@ class SlackMessage:
         re_user = re.compile("<@([^>]+)>")
         user_ids: List[str] = re_user.findall(message)
         users_list = await gather(
-            *(self.workspace.get_user(user_id) for user_id in user_ids)
+            *(self.workspace.users[user_id] for user_id in user_ids)
         )
         users = dict(zip(user_ids, users_list))
 
diff --git a/slack/slack_workspace.py b/slack/slack_workspace.py
index df99f94..5651f4a 100644
--- a/slack/slack_workspace.py
+++ b/slack/slack_workspace.py
@@ -50,6 +50,21 @@ class SlackApi:
         return response
 
 
+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
+
+
 class SlackWorkspace:
     def __init__(self, name: str):
         self.name = name
@@ -57,8 +72,7 @@ class SlackWorkspace:
         self.api = SlackApi(self)
         self.is_connected = False
         self.nick = "TODO"
-        # Maybe make private, so you have to use get_user? Maybe make get_user a getter, though don't know if that's a problem since it's async
-        self.users: Dict[str, Future[SlackUser]] = {}
+        self.users = SlackUsers(self)
         self.conversations: Dict[str, SlackConversation] = {}
 
     async def connect(self):
@@ -85,14 +99,3 @@ class SlackWorkspace:
         # print([c["name"] for c in user_channels])
         self.is_connected = True
         weechat.bar_item_update("input_text")
-
-    async def create_user(self, id: str) -> SlackUser:
-        user = SlackUser(self, id)
-        await user.init()
-        return user
-
-    async def get_user(self, id: str) -> SlackUser:
-        if id in self.users:
-            return await self.users[id]
-        self.users[id] = create_task(self.create_user(id))
-        return await self.users[id]