浏览代码

Refactoring and cleanup

Kirk Trombley 4 年之前
父节点
当前提交
d0963087fb
共有 1 个文件被更改,包括 64 次插入62 次删除
  1. 64 62
      drivers/groupme.py

+ 64 - 62
drivers/groupme.py

@@ -59,69 +59,71 @@ class GroupMeBot(rollbot.Rollbot[GroupMeMessage]):
             attachments=[rollbot.Attachment(d["type"], json.dumps(d)) for d in msg.attachments],
         )
 
+    async def upload_image(self, data: bytes):
+        async with self.context.request.post(
+            "https://image.groupme.com/pictures",
+            headers={
+                "Content-Type": "image/png",
+                "X-Access-Token": groupme_token,
+            },
+            data=data,
+        ) as upload:
+            upload.raise_for_status()
+            return (await upload.json())["payload"]["url"]
+
+    async def post_message(self, bot_id: str, text: str, attachments: list[dict[str, str]]):
+        await self.context.request.post(
+            "https://api.groupme.com/v3/bots/post",
+            json={
+                "bot_id": bot_id,
+                "text": text,
+                "attachments": attachments,
+            },
+            timeout=10,
+        )
+
     async def respond(self, res: rollbot.Response):
-        if res.origin_id == "GROUPME":
-            bot_id = groupme_bots.get(res.channel_id, None)
-            if bot_id is None:
-                self.context.logger.error(f"Unable to respond to group {res.channel_id} in GroupMe")
-                return
-            message = ""
-            attachments = []
-            try:
-                for att in res.attachments:
-                    if att.name == "image":
-                        if isinstance(att.body, bytes):
-                            async with self.context.request.post(
-                                "https://image.groupme.com/pictures",
-                                headers={
-                                    "Content-Type": "image/png",
-                                    "X-Access-Token": groupme_token,
-                                },
-                                data=att.body,
-                            ) as upload:
-                                upload.raise_for_status()
-                                attachments.append(
-                                    {
-                                        "type": "image",
-                                        "url": (await upload.json())["payload"]["url"],
-                                    }
-                                )
-                        else:
-                            attachments.append({"type": "image", "url": att.body})
-            except:
-                self.context.debugging = "".join(traceback.format_exc())
-                message += "Failed to upload one or more attachments!\n"
-                self.context.logger.exception("Failed to upload attachment")
-
-            if res.text is not None:
-                message += res.text
-
-            msgs = []
-            while len(message) > max_msg_len:
-                msgs.append(message[:msg_slice] + split_text)
-                message = message[msg_slice:]
-            msgs.append(message)
-
-            await self.context.request.post(
-                "https://api.groupme.com/v3/bots/post",
-                json={
-                    "bot_id": bot_id,
-                    "text": msgs[0],
-                    "attachments": attachments,
-                },
-                timeout=10,
-            )
-
-            for m in msgs[1:]:
-                await asyncio.sleep(0.1)
-                await self.context.request.post(
-                    "https://api.groupme.com/v3/bots/post",
-                    json={
-                        "bot_id": bot_id,
-                        "text": m,
-                    },
-                    timeout=10,
-                )
+        if res.origin_id != "GROUPME":
+            self.context.logger.error(f"Unable to respond to {res.origin_id}")
+            return
+
+        bot_id = groupme_bots.get(res.channel_id, None)
+        if bot_id is None:
+            self.context.logger.error(f"Unable to respond to group {res.channel_id} in GroupMe")
+            return
+
+        message = ""
+        attachments = []
+        try:
+            for att in res.attachments:
+                if att.name == "image":
+                    if isinstance(att.body, bytes):
+                        attachments.append(
+                            {
+                                "type": "image",
+                                "url": await self.upload_image(att.body),
+                            }
+                        )
+                    else:
+                        attachments.append({"type": "image", "url": att.body})
+        except:
+            self.context.debugging = "".join(traceback.format_exc())
+            message += "Failed to upload one or more attachments!\n"
+            self.context.logger.exception("Failed to upload attachment")
+
+        if res.text is not None:
+            message += res.text
+
+        msgs = []
+        while len(message) > max_msg_len:
+            msgs.append(message[:msg_slice] + split_text)
+            message = message[msg_slice:]
+        msgs.append(message)
+
+        await self.post_message(bot_id, msgs[0], attachments)
+        for m in msgs[1:]:
+            await asyncio.sleep(0.1)
+            await self.post_message(bot_id, m, [])
 
 
 app = FastAPI()