diff --git a/bot/pluralkit/bot.py b/bot/pluralkit/bot.py index be526c8e..6107310d 100644 --- a/bot/pluralkit/bot.py +++ b/bot/pluralkit/bot.py @@ -39,25 +39,29 @@ async def on_message(message): from pluralkit import proxy, utils - # Find and execute command in map - if len(args) > 0 and args[0] in utils.command_map: - subcommand_map = utils.command_map[args[0]] + cmd = None + # Look up commands with subcommands + if len(args) >= 2: + lookup = utils.command_map.get((args[0], args[1]), None) + if lookup: + # Curry with arg slice + cmd = lambda c, m, a: lookup[0](conn, message, args[2:]) + # Look up root commands + if not cmd and len(args) >= 1: + lookup = utils.command_map.get((args[0], None), None) + if lookup: + # Curry with arg slice + cmd = lambda c, m, a: lookup[0](conn, message, args[1:]) - if len(args) >= 2 and args[1] in subcommand_map: - async with client.pool.acquire() as conn: - await subcommand_map[args[1]][0](conn, message, args[2:]) - elif None in subcommand_map: - async with client.pool.acquire() as conn: - await subcommand_map[None][0](conn, message, args[1:]) - elif len(args) >= 2: - embed = discord.Embed() - embed.colour = discord.Colour.dark_red() - embed.description = "Subcommand \"{}\" not found.".format(args[1]) - await client.send_message(message.channel, embed=embed) - else: - # Try doing proxy parsing + # Found anything? run it + if cmd: async with client.pool.acquire() as conn: - await proxy.handle_proxying(conn, message) + await cmd(conn, message, args) + return + + # Try doing proxy parsing + async with client.pool.acquire() as conn: + await proxy.handle_proxying(conn, message) @client.event diff --git a/bot/pluralkit/commands.py b/bot/pluralkit/commands.py index 6b679efb..13affd9a 100644 --- a/bot/pluralkit/commands.py +++ b/bot/pluralkit/commands.py @@ -20,7 +20,7 @@ async def this_system_info(conn, message, args): return True -@command(cmd="pk;system", subcommand="new", usage="[name]", description="Registers a new system to this account.") +@command(cmd="pk;system", subcommand="new", usage="[name]", description="Registers a new system to this account.", basic=True) async def new_system(conn, message, args): system = await db.get_system_by_account(conn, message.author.id) @@ -28,8 +28,8 @@ async def new_system(conn, message, args): return False, "You already have a system registered. To remove your system, use `pk;system remove`, or to unlink your system from this account, use `pk;system unlink`." system_name = None - if len(args) > 2: - system_name = " ".join(args[2:]) + if len(args) > 0: + system_name = " ".join(args) async with conn.transaction(): # TODO: figure out what to do if this errors out on collision on generate_hid @@ -42,7 +42,7 @@ async def new_system(conn, message, args): return True, "System registered! To begin adding members, use `pk;member new `." -@command(cmd="pk;system", subcommand="info", usage="[system]", description="Shows information about a system.") +@command(cmd="pk;system", subcommand="info", usage="[system]", description="Shows information about a system.", basic=True) async def system_info(conn, message, args): if len(args) == 0: # Use sender's system @@ -256,7 +256,7 @@ async def system_fronthistory(conn, message, args): embed.title = "Past switches" return True, embed -@command(cmd="pk;member", subcommand="new", usage="", description="Adds a new member to your system.") +@command(cmd="pk;member", subcommand="new", usage="", description="Adds a new member to your system.", basic=True) async def new_member(conn, message, args): system = await db.get_system_by_account(conn, message.author.id) @@ -276,7 +276,7 @@ async def new_member(conn, message, args): return True, "Member \"{}\" (`{}`) registered!".format(name, hid) -@member_command(cmd="pk;member", subcommand="info", description="Shows information about a system member.", system_only=False) +@member_command(cmd="pk;member", subcommand="info", description="Shows information about a system member.", system_only=False, basic=True) async def member_info(conn, message, member, args): await client.send_message(message.channel, embed=await generate_member_info_card(conn, member)) return True @@ -353,7 +353,7 @@ async def member_remove(conn, message, member, args): return True, "Member removal cancelled." -@member_command(cmd="pk;member", subcommand="avatar", usage="[user|url]", description="Updates a member's avatar. Can be an account mention (which will use that account's avatar), or a link to an image. Leave blank to clear.") +@member_command(cmd="pk;member", subcommand="avatar", usage="[user|url]", description="Updates a member's avatar. Can be an account mention (which will use that account's avatar), or a link to an image. Leave blank to clear.", basic=True) async def member_avatar(conn, message, member, args): if len(args) == 0: avatar_url = None @@ -381,7 +381,7 @@ async def member_avatar(conn, message, member, args): return True, make_default_embed("Avatar set.").set_image(url=avatar_url) -@member_command(cmd="pk;member", subcommand="proxy", usage="[example]", description="Updates a member's proxy settings. Needs an \"example\" proxied message containing the string \"text\" (eg. [text], |text|, etc).") +@member_command(cmd="pk;member", subcommand="proxy", usage="[example]", description="Updates a member's proxy settings. Needs an \"example\" proxied message containing the string \"text\" (eg. [text], |text|, etc).", basic=True) async def member_proxy(conn, message, member, args): if len(args) == 0: prefix, suffix = None, None @@ -455,7 +455,7 @@ async def message_info(conn, message, args): await client.send_message(message.channel, embed=embed) return True -@command(cmd="pk;switch", subcommand=None, usage="", description="Registers a switch and changes the current fronter.") +@command(cmd="pk;switch", subcommand=None, usage="", description="Registers a switch and changes the current fronter.", basic=True) async def switch_member(conn, message, args): if len(args) == 0: return False @@ -508,12 +508,29 @@ def make_help(cmds): cmd, subcmd or "", usage or ""), value=description, inline=False) return embed -@command(cmd="pk;help", subcommand=None, usage="[system|member|message]", description="Shows this help message.") +@command(cmd="pk;help", subcommand=None, usage="[category]", description="Shows this help message.") async def show_help(conn, message, args): - if len(args) > 0 and ("pk;" + args[0]) in command_map: - cmds = ["", ("pk;" + args[0], command_map["pk;" + args[0]])] - else: - cmds = command_map.items() + embed = make_default_embed(None) + embed.title = "PluralKit Help" - await client.send_message(message.channel, embed=make_help(cmds)) - return True + if len(args) == 0: + basics = [] + for (cmd, sub), (_, usage, description, basic) in command_map.items(): + if basic: + basics.append("**{} {} {}** - {}".format(cmd, sub or "", usage or "", description or "")) + embed.add_field(name="Basic commands", value="\n".join(basics)) + + categories = [("system", "System commands"), ("member", "Member commands"), ("switch", "Switching commands")] + categories_lines = ["**pk;help {}** - {}".format(key, desc) for key, desc in categories] + embed.add_field(name="More commands", value="\n".join(categories_lines)) + else: + if args[0] not in ["system", "member", "switch"]: + return False, "Unknown help category." + + cmds = [(k, v) for k, v in command_map.items() if k[0] == "pk;" + args[0]] + lines = [] + for (cmd, sub), (_, usage, description, _) in cmds: + lines.append("**{} {} {}** - {}".format(cmd, sub or "", usage or "", description or "")) + embed.add_field(name="Commands", value="\n".join(lines)) + + return True, embed diff --git a/bot/pluralkit/utils.py b/bot/pluralkit/utils.py index 32664a79..f5ba5279 100644 --- a/bot/pluralkit/utils.py +++ b/bot/pluralkit/utils.py @@ -38,10 +38,11 @@ async def get_system_fuzzy(conn, key) -> asyncpg.Record: if isinstance(key, str) and len(key) == 5: return await db.get_system_by_hid(conn, system_hid=key) - system = parse_mention(key) - - if system: - return system + account = await parse_mention(key) + if account: + system = await db.get_system_by_account(conn, account_id=account.id) + if system: + return system return None @@ -79,7 +80,7 @@ command_map = {} # Second parameter is the message it'll send. If just False, will print usage -def command(cmd, subcommand, usage=None, description=None): +def command(cmd, subcommand, usage=None, description=None, basic=False): def wrap(func): async def wrapper(conn, message, args): res = await func(conn, message, args) @@ -107,12 +108,7 @@ def command(cmd, subcommand, usage=None, description=None): # Success, don't print anything # Put command in map - if cmd not in command_map: - command_map[cmd] = {} - if subcommand not in command_map[cmd]: - command_map[cmd][subcommand] = {} - - command_map[cmd][subcommand] = (wrapper, usage, description) + command_map[(cmd, subcommand)] = (wrapper, usage, description, basic) return wrapper return wrap @@ -121,7 +117,7 @@ def command(cmd, subcommand, usage=None, description=None): # If system_only=False, allows members from other systems by hid -def member_command(cmd, subcommand, usage=None, description=None, system_only=True): +def member_command(cmd, subcommand, usage=None, description=None, system_only=True, basic=False): def wrap(func): async def wrapper(conn, message, args): # Return if no member param @@ -142,20 +138,17 @@ def member_command(cmd, subcommand, usage=None, description=None, system_only=Tr return False, "Can't find member \"{}\".".format(args[0]) return await func(conn, message, member, args[1:]) - return command(cmd=cmd, subcommand=subcommand, usage=" {}".format(usage or ""), description=description)(wrapper) + return command(cmd=cmd, subcommand=subcommand, usage=" {}".format(usage or ""), description=description, basic=basic)(wrapper) return wrap async def generate_system_info_card(conn, system: asyncpg.Record) -> discord.Embed: card = discord.Embed() + card.colour = discord.Colour.blue() if system["name"]: card.title = system["name"] - if system["description"]: - card.add_field(name="Description", - value=system["description"], inline=False) - if system["tag"]: card.add_field(name="Tag", value=system["tag"]) @@ -164,23 +157,20 @@ async def generate_system_info_card(conn, system: asyncpg.Record) -> discord.Emb fronter_val = "{} (for {})".format(current_fronter["name"], humanize.naturaldelta(current_fronter["timestamp"])) card.add_field(name="Current fronter", value=fronter_val) - # Get names of all linked accounts - async def get_name(account_id): - account = await client.get_user_info(account_id) - return "{}#{}".format(account.name, account.discriminator) - - account_name_futures = [] + account_names = [] for account_id in await db.get_linked_accounts(conn, system_id=system["id"]): - account_name_futures.append(get_name(account_id)) - # Run in parallel - account_names = await asyncio.gather(*account_name_futures) - - card.add_field(name="Linked accounts", value=", ".join(account_names)) + account = await client.get_user_info(account_id) + account_names.append("{}#{}".format(account.name, account.discriminator)) + card.add_field(name="Linked accounts", value="\n".join(account_names)) + + if system["description"]: + card.add_field(name="Description", + value=system["description"], inline=False) # Get names of all members member_texts = [] for member in await db.get_all_members(conn, system_id=system["id"]): - member_texts.append("`{}`: {}".format(member["hid"], member["name"])) + member_texts.append("{} (`{}`)".format(member["name"], member["hid"])) if len(member_texts) > 0: card.add_field(name="Members", value="\n".join( @@ -191,8 +181,18 @@ async def generate_system_info_card(conn, system: asyncpg.Record) -> discord.Emb async def generate_member_info_card(conn, member: asyncpg.Record) -> discord.Embed: + system = await db.get_system(conn, system_id=member["system"]) + card = discord.Embed() - card.set_author(name=member["name"], icon_url=member["avatar_url"]) + card.colour = discord.Colour.blue() + + name_and_system = member["name"] + if system["name"]: + name_and_system += " ({})".format(system["name"]) + + card.set_author(name=name_and_system, icon_url=member["avatar_url"] or discord.Embed.Empty) + if member["avatar_url"]: + card.set_thumbnail(url=member["avatar_url"]) if member["color"]: card.colour = int(member["color"], 16) @@ -217,7 +217,7 @@ async def generate_member_info_card(conn, member: asyncpg.Record) -> discord.Emb # Get system name and hid system = await db.get_system(conn, system_id=member["system"]) if system["name"]: - system_value = "`{}`: {}".format(system["hid"], system["name"]) + system_value = "{} (`{}`)".format(system["name"], system["hid"]) else: system_value = "`{}`".format(system["hid"]) card.add_field(name="System", value=system_value, inline=False)