from sqlalchemy import Column, Integer, String from command_system import RollbotResponse, RollbotFailure, as_plugin, ModelBase, pop_arg class CurseBlessScore(ModelBase): __tablename__ = "curse_bless_score" group_id = Column(String, primary_key=True) person_id = Column(Integer, primary_key=True) curses = Column(Integer) blessings = Column(Integer) def get_score(db, group_id, name): person_id = name.strip().lower() 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 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.") @as_plugin def curse(db, msg): # TODO might be nice to add subcommands to this later name, args = pop_arg(msg.raw_args) if not name.startswith("!"): score = get_score(db, msg.group_id, name) # 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) # strip off the '!' subc = name[1:].strip() if len(subc) == 0: # 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) @as_plugin def bless(db, msg): name, _ = pop_arg(msg.raw_args) score = get_score(db, msg.group_id, name) score.blessings = score.blessings + 1 return get_response(msg, name, score)