浏览代码

Refactoring curse/bless

Kirk Trombley 4 年之前
父节点
当前提交
fc702d27ce
共有 2 个文件被更改,包括 75 次插入87 次删除
  1. 4 0
      config/config.toml
  2. 71 87
      src/plugins/curse.py

+ 4 - 0
config/config.toml

@@ -13,6 +13,10 @@ fs = "f"
 hg = "hangguy"
 cookie = "fortune"
 wl = "watchlist"
+bless = "curse"
+blurse = "curse"
+uncurse = "curse"
+unbless = "curse"
 
 [responses]
 info = """

+ 71 - 87
src/plugins/curse.py

@@ -1,8 +1,9 @@
 import logging
+import random
 
 from sqlalchemy import Column, Integer, String
 
-from rollbot import RollbotResponse, RollbotFailure, RollbotPlugin, as_plugin, ModelBase, pop_arg
+from rollbot import RollbotFailure, as_plugin, ModelBase, pop_arg
 
 
 class CurseBlessScore(ModelBase):
@@ -29,10 +30,6 @@ def fmt_times(n):
     return str(n) + (" time" if n == 1 else " times")
 
 
-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()
@@ -41,33 +38,33 @@ def curse_agg(db, msg, args):
     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.")
+    return 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)
     if name is None:
-        return RollbotResponse(msg, txt=f"Sorry! I need to know who's records you're trying to clear!")
+        return f"Sorry! I need to know who's records you're trying to clear!"
     person_id = name.strip().lower()
     if not msg.from_admin:
-        return RollbotResponse(msg, failure=RollbotFailure.PERMISSIONS)
+        return 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!")
+        return 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)}!")
+    return 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)
     if name is None:
-        return RollbotResponse(msg, txt=f"Sorry! I need to know who's records you're trying to clear!")
+        return f"Sorry! I need to know who's records you're trying to clear!"
     person_id = name.strip().lower()
     if not msg.from_admin:
-        return RollbotResponse(msg, failure=RollbotFailure.PERMISSIONS)
+        return RollbotFailure.PERMISSIONS
     cnt = 0
     oc = 0
     ob = 0
@@ -78,20 +75,37 @@ def curse_clearall(db, msg, args):
         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!")
+        return f"Sorry! I don't know who {name} is!"
+    return 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!"
+
+
+TOP_MAP = {
+    "curse": CurseBlessScore.curses,
+    "bless": CurseBlessScore.blessings,
+}
 
 
 def curse_top(db, msg, args):
-    rankings = "\n".join(f"{i + 1}. {cbs.person_id.capitalize()} was cursed {fmt_times(cbs.curses)}"
+    if msg.command == "curse":
+        target = CurseBlessScore.curses
+        text = "cursed"
+        extract = lambda cbs: cbs.curses
+    elif msg.command == "bless":
+        target = CurseBlessScore.blessings
+        text = "blessed"
+        extract = lambda cbs: cbs.blessings
+    else:
+        return "Sorry! Right now !top can only be used with curse and bless."
+    
+    rankings = "\n".join(f"{i + 1}. {cbs.person_id.capitalize()} was {text} {fmt_times(extract(cbs))}"
         for i, cbs in enumerate(db.query(CurseBlessScore
             ).filter(CurseBlessScore.group_id == msg.group_id
-            ).filter(CurseBlessScore.curses != 0
-            ).order_by(CurseBlessScore.curses.desc()
+            ).filter(target != 0
+            ).order_by(target.desc()
             ).limit(10)))
     if len(rankings) == 0:
-        return RollbotResponse(msg, txt="Sorry! I don't have enough curse history in this chat to give statistics.")
-    return RollbotResponse(msg, txt=f"The worst offenders of this chat:\n{rankings}")
+        return "Sorry! I don't have enough curse history in this chat to give statistics."
+    return f"The worst offenders of this chat:\n{rankings}"
 
 
 SUBC_MAP = {
@@ -102,82 +116,52 @@ SUBC_MAP = {
 }
 
 
-@as_plugin
-def curse(db, msg, subc, bot, ban_list_cfg: "curse.banlist", no_list_cfg: "curse.nolist"):
-    if subc is None:
-        name, _ = pop_arg(msg.raw_args)
-        if name is None:
-            return "Sorry - you need to provide the name of someone to curse!"
-        person_id = name.strip().lower()
-        if is_banned(msg, person_id, ban_list_cfg):
-            if msg.sender_id in no_list_cfg:
-                return "No!!!!!!!!!!"
-            return "Hey! You aren't allowed to affect that person's score!"
-        score = get_score(db, msg.group_id, person_id)
-        # Note we do this instead of += b/c that can create race conditions in sqlalchemy
-        score.curses = score.curses + 1
-        return get_response(msg, name, score)
-
-    return SUBC_MAP.get(subc.command, lambda *_: RollbotResponse(msg, failure=RollbotFailure.INVALID_SUBCOMMAND))(db, msg, subc.raw_args)
-
-
-class Bless(RollbotPlugin):
-    def __init__(self, bot, logger=logging.getLogger(__name__)):
-        RollbotPlugin.__init__(self, "bless", bot, logger=logger)
-        self.ban_list = bot.config.get("curse.banlist")
+def on_curse(score):
+    # Note we do (x := x + 1) instead of += b/c that can create race conditions in sqlalchemy
+    score.curses = score.curses + 1
+    return score
 
-    def do_score_change(self, score):
-        score.blessings = score.blessings + 1
 
-    def on_command(self, db, msg):
-        name, _ = pop_arg(msg.raw_args)
-        if name.startswith("!"):
-            return RollbotResponse(msg, txt="Sorry! Subcommands have to go on !curse for now - this will be fixed in the future!")
-        person_id = name.strip().lower()
-        if is_banned(msg, person_id, self.ban_list):
-            self.bot.manually_post_message("Hey! You aren't allowed to affect that person's score! And cheaters never propser!", msg.group_id)
-            return RollbotResponse(msg, txt=f"!curse {self.ban_list[msg.sender_id][0]} for trying to bless someone they can't!")
-        score = get_score(db, msg.group_id, person_id)
-        self.do_score_change(score)
-        return get_response(msg, name, score)
+def on_bless(score):
+    score.blessings = score.blessings + 1
+    return score
 
 
-class Blurse(Bless):
-    def __init__(self, bot, logger=logging.getLogger(__name__)):
-        RollbotPlugin.__init__(self, "blurse", bot, logger=logger)
+def on_uncurse(score):
+    score.curses = score.curses - 1
+    return score
 
-    def do_score_change(self, score):
-        score.blessings = score.blessings + 1
-        score.curses = score.curses + 1
 
+def on_unbless(score):
+    score.blessings = score.blessings - 1
+    return score
 
-@as_plugin
-def unbless(db, msg, ban_list_cfg: "curse.banlist"):
-    name, _ = pop_arg(msg.raw_args)
-    if name.startswith("!"):
-        return "Sorry! Subcommands have to go on !curse for now - this will be fixed in the future!"
-    if not msg.from_admin:
-        return RollbotResponse(msg, failure=RollbotFailure.PERMISSIONS, debugging={ "explain": "For now, only admins can unbless or uncurse people!" })
-    person_id = name.strip().lower()
-    if is_banned(msg, person_id, ban_list_cfg):
-        return "Hey! You aren't allowed to affect that person's score!"
 
-    score = get_score(db, msg.group_id, person_id)
-    score.blessings = score.blessings - 1
-    return get_response(msg, name, score)
+CMD_MAP = {
+    "curse": on_curse,
+    "bless": on_bless,
+    "blurse": lambda score: on_bless(on_curse(score)),
+    "uncurse": on_uncurse,
+    "unbless": on_unbless,
+}
 
 
 @as_plugin
-def uncurse(db, msg, ban_list_cfg: "curse.banlist"):
-    name, _ = pop_arg(msg.raw_args)
-    if name.startswith("!"):
-        return "Sorry! Subcommands have to go on !curse for now - this will be fixed in the future!"
-    if not msg.from_admin:
-        return RollbotResponse(msg, failure=RollbotFailure.PERMISSIONS, debugging={ "explain": "For now, only admins can unbless or uncurse people!" })
-    person_id = name.strip().lower()
-    if is_banned(msg, person_id, ban_list_cfg):
-        return "Hey! You aren't allowed to affect that person's score!"
-
-    score = get_score(db, msg.group_id, person_id)
-    score.curses = score.curses - 1
-    return get_response(msg, name, score)
+def curse(db, msg, subc, bot, ban_list_cfg: "curse.banlist", no_list_cfg: "curse.nolist"):
+    if subc is None:
+        if not msg.from_admin and msg.command.startswith("un"):
+            return RollbotFailure.PERMISSIONS.with_reason("For now, only admins can unbless or uncurse people!")
+        if msg.raw_args is None:
+            return f"Sorry - you need to provide the name of someone to {msg.command}!"
+        person_id = msg.raw_args.strip().lower().replace(" ", "_")
+        if is_banned(msg, person_id, ban_list_cfg):
+            if msg.sender_id in no_list_cfg:
+                return "No!!!!!!!!!!"
+            if random.randint(0, 1) == 0:
+                bot.manually_post_message("Hey! You aren't allowed to affect that person's score! And cheaters never propser!", msg.group_id)
+                return f"!curse {ban_list_cfg[msg.sender_id][0]}"
+            return "Hey! You aren't allowed to affect that person's score!"
+        score = get_score(db, msg.group_id, person_id)
+        CMD_MAP[msg.command](score)
+        return f"It is done! {msg.raw_args} has been cursed {fmt_times(score.curses)} and blessed {fmt_times(score.blessings)} in this chat."
+    return SUBC_MAP.get(subc.command, lambda *_: RollbotFailure.INVALID_SUBCOMMAND)(db, msg, subc.raw_args)