浏览代码

Add scp command

Kirk Trombley 4 年之前
父节点
当前提交
65664aca17
共有 2 个文件被更改,包括 41 次插入1 次删除
  1. 40 1
      src/commands/querying.py
  2. 1 0
      src/setup.py

+ 40 - 1
src/commands/querying.py

@@ -1,7 +1,10 @@
 from logging import Logger
+import asyncio
+
+from bs4 import BeautifulSoup
 
 from rollbot import as_command, RollbotFailure, Attachment
-from rollbot.injection import Request, Args
+from rollbot.injection import Request, Args, Arg
 
 
 @as_command
@@ -40,3 +43,39 @@ async def shield(req: Request, blazon: Args, logger: Logger):
     except:
         logger.exception("Failed !shield")
         RollbotFailure.SERVICE_DOWN.raise_exc(detail="Could not reach DrawShield")
+
+
+@as_command
+async def scp(
+    number: Arg(0, convert=int, fail_msg="Could not parse argument {} into integer"), req: Request
+):
+    page_url = f"http://www.scp-wiki.net/scp-{number:03d}"
+    series_url = "http://www.scp-wiki.net/scp-series"
+    series = (number // 1000) + 1
+    if series != 1:
+        series_url += "-" + str(series)
+
+    async def get_obj_class():
+        async with req.get(page_url) as page:
+            page_html = BeautifulSoup(await page.text(), "html.parser")
+
+        for p in page_html.find_all("p"):
+            if "Object Class" in p.text:
+                return p.text
+
+        return "Error retrieving object class!"
+
+    query = f"SCP-{number:03d}"
+
+    async def get_title():
+        async with req.get(series_url) as series:
+            series_html = BeautifulSoup(await series.text(), "html.parser")
+
+        for li in series_html.find_all("li"):
+            if query in li.text:
+                return li.text.split("-", 2)[-1].strip()
+
+        return "Error retrieving title!"
+
+    obj_class, title = await asyncio.gather(get_obj_class(), get_title())
+    return f"Item # {number}\n{obj_class}\n{title}\n{page_url}"

+ 1 - 0
src/setup.py

@@ -12,5 +12,6 @@ setup(
     install_requires=[
         "rollbot",
         "dice",
+        "bs4",
     ],
 )