瀏覽代碼

Refactoring subcommand dispatch for curse

Kirk Trombley 6 年之前
父節點
當前提交
390313bf11
共有 1 個文件被更改,包括 53 次插入38 次删除
  1. 53 38
      src/plugins/curse.py

+ 53 - 38
src/plugins/curse.py

@@ -29,6 +29,58 @@ def get_response(msg, name, score):
     return RollbotResponse(msg, txt=f"It is done! {name} has been cursed {fmt_times(score.curses)} and blessed {fmt_times(score.blessings)} in this chat.")
 
 
+def curse_agg(db, msg, args):
+    name, _ = pop_arg(args)
+    person_id = name.strip().lower()
+    curses = 0
+    blessings = 0
+    for cbs in db.query(CurseBlessScore).filter(CurseBlessScore.person_id == person_id):
+        curses += cbs.curses
+        blessings += cbs.blessings
+    return RollbotResponse(msg, txt=f"From my records, {name} has been cursed {fmt_times(curses)} and blessed {fmt_times(blessings)} across all chats.")
+
+
+def curse_clear(db, msg, args):
+    name, _ = pop_arg(args)
+    person_id = name.strip().lower()
+    if not msg.from_admin:
+        return RollbotResponse(msg, failure=RollbotFailure.PERMISSIONS)
+    cbs = db.query(CurseBlessScore).get((msg.group_id, person_id))
+    if cbs is None:
+        return RollbotResponse(msg, txt=f"Sorry! I don't know who {name} is!")
+    oc = cbs.curses
+    ob = cbs.blessings
+    cbs.curses = 0
+    cbs.blessings = 0
+    return RollbotResponse(msg, txt=f"Done! I have cleared all records of {name} in this chat, who was previously cursed {fmt_times(oc)} and blessed {fmt_times(ob)}!")
+
+
+def curse_clearall(db, msg, args):
+    name, _ = pop_arg(args)
+    person_id = name.strip().lower()
+    if not msg.from_admin:
+        return RollbotResponse(msg, failure=RollbotFailure.PERMISSIONS)
+    cnt = 0
+    oc = 0
+    ob = 0
+    for cbs in db.query(CurseBlessScore).filter(CurseBlessScore.person_id == person_id):
+        oc += cbs.curses
+        ob += cbs.blessings
+        cbs.curses = 0
+        cbs.blessings = 0
+        cnt += 1
+    if cnt == 0:
+        return RollbotResponse(msg, txt=f"Sorry! I don't know who {name} is!")
+    return RollbotResponse(msg, txt=f"Done! I have cleared all records of {name} in all chats, who was previously cursed {fmt_times(oc)} and blessed {fmt_times(ob)} in total!")
+
+
+SUBC_MAP = {
+    "aggregate": curse_agg,
+    "clear": curse_clear,
+    "clearall": curse_clearall,
+}
+
+
 @as_plugin
 def curse(db, msg):
     # TODO might be nice to add subcommands to this later
@@ -45,45 +97,8 @@ def curse(db, msg):
         # handle the case of spaces between ! and subcommand
         subc, args = pop_arg(args)
     subc = subc.lower()
-    name, _ = pop_arg(args)
-    person_id = name.strip().lower()
-
-    if subc == "aggregate":
-        curses = 0
-        blessings = 0
-        for cbs in db.query(CurseBlessScore).filter(CurseBlessScore.person_id == person_id):
-            curses += cbs.curses
-            blessings += cbs.blessings
-        return RollbotResponse(msg, txt=f"From my records, {name} has been cursed {fmt_times(curses)} and blessed {fmt_times(blessings)} across all chats.")
-    elif subc == "clear":
-        if not msg.from_admin:
-            return RollbotResponse(msg, failure=RollbotFailure.PERMISSIONS)
-        cbs = db.query(CurseBlessScore).get((msg.group_id, person_id))
-        if cbs:
-            oc = cbs.curses
-            ob = cbs.blessings
-            cbs.curses = 0
-            cbs.blessings = 0
-            return RollbotResponse(msg, txt=f"Done! I have cleared all records of {name} in this chat, who was previously cursed {fmt_times(oc)} and blessed {fmt_times(ob)}!")
-        return RollbotResponse(msg, txt=f"Sorry! I don't know who {name} is!")
-    elif subc == "clearall":
-        if not msg.from_admin:
-            return RollbotResponse(msg, failure=RollbotFailure.PERMISSIONS)
-        cnt = 0
-        oc = 0
-        ob = 0
-        for cbs in db.query(CurseBlessScore).filter(CurseBlessScore.person_id == person_id):
-            oc += cbs.curses
-            ob += cbs.blessings
-            cbs.curses = 0
-            cbs.blessings = 0
-            cnt += 1
-        if cnt == 0:
-            return RollbotResponse(msg, txt=f"Sorry! I don't know who {name} is!")
-        return RollbotResponse(msg, txt=f"Done! I have cleared all records of {name} in all chats, who was previously cursed {fmt_times(oc)} and blessed {fmt_times(ob)} in total!")
-    else:
-        return RollbotResponse(msg, failure=RollbotFailure.INVALID_SUBCOMMAND)
 
+    return SUBC_MAP.get(subc, lambda *a: RollbotResponse(msg, failure=RollbotFailure.INVALID_SUBCOMMAND))(db, msg, args)
 
 
 @as_plugin