rollbot3/src/plugins/watchlist.py

import pickle
import random

from command_system 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(msg, data: Watchlist):
    if msg.raw_args is not None and msg.raw_args.startswith("!"):
        subc, args = pop_arg(msg.raw_args)
        subc = subc[1:].lower()
        subc_fn = subcommands.get(subc, lambda *_: f"Sorry, subcommand must be one of {', '.join(subcommands)}, you used: {subc}")
        return subc_fn(args, data)
    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)