import pickle import random from rollbot import as_plugin, as_group_singleton, pop_arg @as_group_singleton class Watchlist: items: "binary" def get_items(self): if self.items is None: return [] return pickle.loads(self.items) def set_items(self, new_items): new_items = new_items or [] self.items = pickle.dumps(new_items) def view(args, data): items = data.get_items() if len(items) == 0: return "This chat hasn't added any items to the watch list!" return "\n".join(f"#{i + 1} - {x}" for i, x in enumerate(items)) def add(args, data): to_add = [x.strip() for x in args.split(";") if len(x) > 0 and not x.isspace()] if len(to_add) == 0: return "Sorry, doesn't look like you provided any valid items! You can also add multiple items using ; as a delimiter." items = data.get_items() new_items = items + to_add data.set_items(new_items) item_list = "\n".join(f"#{i + len(items) + 1} - {x}" for i, x in enumerate(to_add)) return f"Added {len(to_add)} item(s):\n" + item_list def pick(args, data): items = data.get_items() if len(items) == 0: return "You have to add items before I can pick one!" ind = random.randint(0, len(items) - 1) chosen = items[ind] return f"I have selected #{ind + 1} - {chosen}! If you want to remove this from the list now, you can use !wlr {ind + 1}" def remove(args, data): ind, _ = pop_arg(args) try: ind = int(ind) - 1 except ValueError: return f"Sorry, {ind} is not a valid index." items = data.get_items() if ind < 0 or ind >= len(items): return f"Sorry, {ind + 1} is not valid for the list, which ranges from 1 to {len(items)}" removed = items[ind] del items[ind] data.set_items(items) return f"Done! I have removed {removed} from the list.\nNote this might change all the other numbers, so do a !wlv before you do more removes!" HELP_MSG=""" This command lets you track a list of things to watch Useful subcommands are: !watchlist !help - show this message !watchlist !add OR !wla - to add new things to the list, delimited by ;'s. !watchlist !view OR !wlv - to view the current list. !watchlist !pick OR !wlp - to pick something random from the list. !watchlist !remove OR !wlr - to remove something from the list by number. """.strip() subcommands = { "help": lambda *_: HELP_MSG, "view": view, "add": add, "pick": pick, "remove": remove, } @as_plugin def watchlist(subc, data: Watchlist): if subc: subc_fn = subcommands.get(subc.command, None) if subc_fn: return subc_fn(subc.raw_args, data) return f"Sorry, subcommand must be one of {', '.join(subcommands)}, you used: {subc.command}" return subcommands["help"]() # default to help message # some quick aliases for convenience, !wl is aliased in the config file instead @as_plugin def wlp(msg, data: Watchlist): return pick(msg.raw_args, data) @as_plugin def wlv(msg, data: Watchlist): return view(msg.raw_args, data) @as_plugin def wlr(msg, data: Watchlist): return remove(msg.raw_args, data) @as_plugin def wla(msg, data: Watchlist): return add(msg.raw_args, data)