Fix bug in injector, refactor curse function, add additional Query options

lib/rollbot/plugins/injection.py

@@ -78,17 +78,32 @@ class Query(ArgConverter):
     def __init__(self, model_cls):
         super().__init__(self._convert)
         self.model_cls = model_cls
-        self.filters = []
+        self._filters = []
+        self._order_by = None
+        self._limit = None
 
     def filter(self, fn):
-        self.filters.append(fn)
+        self._filters.append(fn)
+        return self
+
+    def order_by(self, fn):
+        self._order_by = fn
+        return self
+
+    def limit(self, n):
+        self._limit = n
         return self
 
     def _convert(self, cmd, db, msg):
         query = db.query(self.model_cls)
-        for fn in self.filters:
+        for fn in self._filters:
             query = query.filter(_run_converter_function(fn, cmd, db, msg))
-        return query.all()
+        if self._order_by is not None:
+            query = query.order_by(_run_converter_function(self._order_by, cmd, db, msg))
+        if self._limit is None:
+            return query.all()
+        else:
+            return query.limit(self._limit)
 
 
 class Lazy(ArgConverter):
@@ -117,7 +132,7 @@ class _Executor:
         self.msg = msg
 
     def run_with_deps(self, fn):
-        return _run_converter_function(fn, cmd, db, msg)
+        return _run_converter_function(fn, self.cmd, self.db, self.msg)
         
 
 Injector = ArgConverter(_Executor)

src/plugins/curse.py

@@ -19,18 +19,6 @@ def get_person_id(args):
     return args.strip().lower().replace(" ", "_")
 
 
-def get_score(db, group_id, person_id):
-    score = db.query(CurseBlessScore).get((group_id, person_id))
-    if not score:
-        score = CurseBlessScore(group_id=group_id, person_id=person_id, curses=0, blessings=0)
-        db.add(score)
-    return score
-
-
-def is_banned(msg, person_id, ban_list):
-    return msg.sender_id in ban_list and person_id in ban_list[msg.sender_id]
-
-
 def fmt_times(n):
     return str(n) + (" time" if n == 1 else " times")
 
@@ -83,12 +71,6 @@ def curse_clearall(db: Database, msg: Message, args: Subcommand.ArgString):
     return f"Done! I have cleared all records of {args} 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: Database, msg: Message, args: Subcommand.ArgString):
     if msg.command == "curse":
         target = CurseBlessScore.curses
@@ -150,31 +132,41 @@ CMD_MAP = {
 }
 
 
+def curse_command(db: Database, 
+                  msg: Message,
+                  bot: Bot, 
+                  ban_list: Config("curse.banlist"), 
+                  no_list: Config("curse.nolist"),
+                  inject: Injector):
+    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 = get_person_id(msg.raw_args)
+    if msg.sender_id in ban_list and person_id in ban_list[msg.sender_id]:
+        if msg.sender_id in no_list:
+            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[msg.sender_id][0]}"
+        return "Hey! You aren't allowed to affect that person's score!"
+
+    score = db.query(CurseBlessScore).get((msg.group_id, person_id))
+    if score is None:
+        score = CurseBlessScore(group_id=msg.group_id, person_id=person_id, curses=0, blessings=0)
+        db.add(score)
+    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."
+
+
 @as_plugin
-def curse(db: Database, 
-          msg: Message, 
-          subc: Subcommand, 
-          bot: Bot, 
-          ban_list: Config("curse.banlist"), 
-          no_list: Config("curse.nolist"),
-          inject: Injector):
+def curse(subc: Subcommand, injector: Injector):
     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 = get_person_id(msg.raw_args)
-        if is_banned(msg, person_id, ban_list):
-            if msg.sender_id in no_list:
-                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[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 injector.run_with_deps(curse_command)
     try:
-        return inject.run_with_deps(SUBC_MAP[subc.command])
+        return injector.run_with_deps(SUBC_MAP[subc.command])
     except KeyError:
-        RollbotFailure.INVALID_SUBCOMMAND.raise_exc()
+        return RollbotFailure.INVALID_SUBCOMMAND