From 902c14ef654966ae7f1ad6832476114311080bf6 Mon Sep 17 00:00:00 2001 From: Ske Date: Thu, 12 Jul 2018 00:49:02 +0200 Subject: [PATCH] Format code --- bot/pluralkit/__init__.py | 2 +- bot/pluralkit/bot.py | 15 ++++++-- bot/pluralkit/commands.py | 80 +++++++++++++++++++++++++-------------- bot/pluralkit/db.py | 59 +++++++++++++++++++++++------ bot/pluralkit/proxy.py | 22 +++++++---- bot/pluralkit/utils.py | 44 ++++++++++++++------- 6 files changed, 157 insertions(+), 65 deletions(-) diff --git a/bot/pluralkit/__init__.py b/bot/pluralkit/__init__.py index 37d28ec9..b2a23fcc 100644 --- a/bot/pluralkit/__init__.py +++ b/bot/pluralkit/__init__.py @@ -1 +1 @@ -from . import commands, db, proxy \ No newline at end of file +from . import commands, db, proxy diff --git a/bot/pluralkit/bot.py b/bot/pluralkit/bot.py index fcc1ecf6..be526c8e 100644 --- a/bot/pluralkit/bot.py +++ b/bot/pluralkit/bot.py @@ -12,23 +12,28 @@ logger.setLevel(logging.DEBUG) client = discord.Client() + @client.event async def on_error(evt, *args, **kwargs): - logger.exception("Error while handling event {} with arguments {}:".format(evt, args)) + logger.exception( + "Error while handling event {} with arguments {}:".format(evt, args)) + @client.event async def on_ready(): # Print status info logger.info("Connected to Discord.") - logger.info("Account: {}#{}".format(client.user.name, client.user.discriminator)) + logger.info("Account: {}#{}".format( + client.user.name, client.user.discriminator)) logger.info("User ID: {}".format(client.user.id)) + @client.event async def on_message(message): # Ignore bot messages if message.author.bot: return - + # Split into args. shlex sucks so we don't bother with quotes args = message.content.split(" ") @@ -54,6 +59,7 @@ async def on_message(message): async with client.pool.acquire() as conn: await proxy.handle_proxying(conn, message) + @client.event async def on_reaction_add(reaction, user): from pluralkit import proxy @@ -62,6 +68,7 @@ async def on_reaction_add(reaction, user): async with client.pool.acquire() as conn: await proxy.handle_reaction(conn, reaction, user) + async def run(): from pluralkit import db try: @@ -79,4 +86,4 @@ async def run(): await client.start(os.environ["TOKEN"]) finally: logger.info("Logging out from Discord...") - await client.logout() \ No newline at end of file + await client.logout() diff --git a/bot/pluralkit/commands.py b/bot/pluralkit/commands.py index a7fc9af2..28f5274c 100644 --- a/bot/pluralkit/commands.py +++ b/bot/pluralkit/commands.py @@ -8,6 +8,7 @@ from pluralkit import db from pluralkit.bot import client, logger from pluralkit.utils import command, generate_hid, generate_member_info_card, generate_system_info_card, member_command, parse_mention, text_input, get_system_fuzzy, get_member_fuzzy, command_map + @command(cmd="pk;system", subcommand=None, description="Shows information about your system.") async def this_system_info(conn, message, args): system = await db.get_system_by_account(conn, message.author.id) @@ -18,6 +19,7 @@ async def this_system_info(conn, message, args): await client.send_message(message.channel, embed=await generate_system_info_card(conn, system)) return True + @command(cmd="pk;system", subcommand="new", usage="[name]", description="Registers a new system to this account.") async def new_system(conn, message, args): system = await db.get_system_by_account(conn, message.author.id) @@ -39,12 +41,13 @@ async def new_system(conn, message, args): await db.link_account(conn, system_id=system["id"], account_id=message.author.id) 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.") async def system_info(conn, message, args): if len(args) == 0: # Use sender's system system = await db.get_system_by_account(conn, message.author.id) - + if system is None: return False, "No system is registered to this account." else: @@ -53,14 +56,15 @@ async def system_info(conn, message, args): if system is None: return False, "Unable to find system \"{}\".".format(args[0]) - + await client.send_message(message.channel, embed=await generate_system_info_card(conn, system)) return True + @command(cmd="pk;system", subcommand="name", usage="[name]", description="Renames your system. Leave blank to clear.") async def system_name(conn, message, args): system = await db.get_system_by_account(conn, message.author.id) - + if system is None: return False, "No system is registered to this account." @@ -73,10 +77,11 @@ async def system_name(conn, message, args): await db.update_system_field(conn, system_id=system["id"], field="name", value=new_name) return True, "Name updated to {}.".format(new_name) if new_name else "Name cleared." + @command(cmd="pk;system", subcommand="description", usage="[clear]", description="Updates your system description. Add \"clear\" to clear.") async def system_description(conn, message, args): system = await db.get_system_by_account(conn, message.author.id) - + if system is None: return False, "No system is registered to this account." @@ -93,10 +98,11 @@ async def system_description(conn, message, args): await db.update_system_field(conn, system_id=system["id"], field="description", value=new_description) return True, "Description set." if new_description else "Description cleared." + @command(cmd="pk;system", subcommand="tag", usage="[tag]", description="Updates your system tag. Leave blank to clear.") async def system_tag(conn, message, args): system = await db.get_system_by_account(conn, message.author.id) - + if system is None: return False, "No system is registered to this account." @@ -110,24 +116,27 @@ async def system_tag(conn, message, args): members_exceeding = await db.get_members_exceeding(conn, system_id=system["id"], length=max_length - len(tag)) if len(members_exceeding) > 0: # If so, error out and warn - member_names = ", ".join([member["name"] for member in members_exceeding]) - logger.debug("Members exceeding combined length with tag '{}': {}".format(tag, member_names)) + member_names = ", ".join([member["name"] + for member in members_exceeding]) + logger.debug("Members exceeding combined length with tag '{}': {}".format( + tag, member_names)) return False, "The maximum length of a name plus the system tag is 32 characters. The following members would exceed the limit: {}. Please reduce the length of the tag, or rename the members.".format(member_names) async with conn.transaction(): await db.update_system_field(conn, system_id=system["id"], field="tag", value=tag) - + return True, "Tag updated to {}.".format(tag) if tag else "Tag cleared." + @command(cmd="pk;system", subcommand="remove", description="Removes your system ***permanently***.") async def system_remove(conn, message, args): system = await db.get_system_by_account(conn, message.author.id) - + if system is None: return False, "No system is registered to this account." await client.send_message(message.channel, "Are you sure you want to remove your system? If so, reply to this message with the system's ID (`{}`).".format(system["hid"])) - + msg = await client.wait_for_message(author=message.author, channel=message.channel) if msg.content == system["hid"]: await db.remove_system(conn, system_id=system["id"]) @@ -139,7 +148,7 @@ async def system_remove(conn, message, args): @command(cmd="pk;system", subcommand="link", usage="", description="Links another account to your system.") async def system_link(conn, message, args): system = await db.get_system_by_account(conn, message.author.id) - + if system is None: return False, "No system is registered to this account." @@ -176,7 +185,7 @@ async def system_link(conn, message, args): @command(cmd="pk;system", subcommand="unlink", description="Unlinks your system from this account. There must be at least one other account linked.") async def system_unlink(conn, message, args): system = await db.get_system_by_account(conn, message.author.id) - + if system is None: return False, "No system is registered to this account." @@ -184,15 +193,16 @@ async def system_unlink(conn, message, args): linked_accounts = await db.get_linked_accounts(conn, system_id=system["id"]) if len(linked_accounts) == 1: return False, "This is the only account on your system, so you can't unlink it." - + async with conn.transaction(): await db.unlink_account(conn, system_id=system["id"], account_id=message.author.id) return True, "Account unlinked." + @command(cmd="pk;member", subcommand="new", usage="", description="Adds a new member to your system.") async def new_member(conn, message, args): system = await db.get_system_by_account(conn, message.author.id) - + if system is None: return False, "No system is registered to this account." @@ -208,11 +218,13 @@ async def new_member(conn, message, args): await db.create_member(conn, system_id=system["id"], member_name=name, member_hid=hid) return True, "Member \"{}\" (`{}`) registered!".format(name, hid) + @member_command(cmd="pk;member", subcommand="info", description="Shows information about a system member.", system_only=False) async def member_info(conn, message, member, args): await client.send_message(message.channel, embed=await generate_member_info_card(conn, member)) return True + @member_command(cmd="pk;member", subcommand="color", usage="[color]", description="Updates a member's associated color. Leave blank to clear.") async def member_color(conn, message, member, args): if len(args) == 0: @@ -227,7 +239,8 @@ async def member_color(conn, message, member, args): async with conn.transaction(): await db.update_member_field(conn, member_id=member["id"], field="color", value=color) return True, "Color updated to #{}.".format(color) if color else "Color cleared." - + + @member_command(cmd="pk;member", subcommand="pronouns", usage="[pronouns]", description="Updates a member's pronouns. Leave blank to clear.") async def member_pronouns(conn, message, member, args): if len(args) == 0: @@ -239,12 +252,13 @@ async def member_pronouns(conn, message, member, args): await db.update_member_field(conn, member_id=member["id"], field="pronouns", value=pronouns) return True, "Pronouns set to {}".format(pronouns) if pronouns else "Pronouns cleared." + @member_command(cmd="pk;member", subcommand="birthdate", usage="[birthdate]", description="Updates a member's birthdate. Must be in ISO-8601 format (eg. 1999-07-25). Leave blank to clear.") async def member_birthday(conn, message, member, args): if len(args) == 0: new_date = None else: - # Parse date + # Parse date try: new_date = datetime.strptime(args[0], "%Y-%m-%d").date() except ValueError: @@ -254,6 +268,7 @@ async def member_birthday(conn, message, member, args): await db.update_member_field(conn, member_id=member["id"], field="birthday", value=new_date) return True, "Birthdate set to {}".format(new_date) if new_date else "Birthdate cleared." + @member_command(cmd="pk;member", subcommand="description", description="Updates a member's description. Add \"clear\" to clear.") async def member_description(conn, message, member, args): if len(args) > 0 and args[0] == "clear": @@ -267,11 +282,12 @@ async def member_description(conn, message, member, args): async with conn.transaction(): await db.update_member_field(conn, member_id=member["id"], field="description", value=new_description) return True, "Description set." if new_description else "Description cleared." - + + @member_command(cmd="pk;member", subcommand="remove", description="Removes a member from your system.") async def member_remove(conn, message, member, args): await client.send_message(message.channel, "Are you sure you want to remove {}? If so, reply to this message with the member's name.".format(member["name"])) - + msg = await client.wait_for_message(author=message.author, channel=message.channel) if msg.content == member["name"]: await db.delete_member(conn, member_id=member["id"]) @@ -279,6 +295,7 @@ async def member_remove(conn, message, member, args): else: 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.") async def member_avatar(conn, message, member, args): if len(args) == 0: @@ -296,11 +313,11 @@ async def member_avatar(conn, message, member, args): avatar_url = args[0] else: return False, "Invalid URL." - + async with conn.transaction(): await db.update_member_field(conn, member_id=member["id"], field="avatar_url", value=avatar_url) return True, "Avatar set." if avatar_url else "Avatar cleared." - + @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).") async def member_proxy(conn, message, member, args): @@ -318,7 +335,8 @@ async def member_proxy(conn, message, member, args): # Extract prefix and suffix prefix = example[:example.index("text")].strip() suffix = example[example.index("text")+4:].strip() - logger.debug("Matched prefix '{}' and suffix '{}'".format(prefix, suffix)) + logger.debug( + "Matched prefix '{}' and suffix '{}'".format(prefix, suffix)) # DB stores empty strings as None, make that work if not prefix: @@ -331,6 +349,7 @@ async def member_proxy(conn, message, member, args): await db.update_member_field(conn, member_id=member["id"], field="suffix", value=suffix) return True, "Proxy settings updated." if prefix or suffix else "Proxy settings cleared." + @command(cmd="pk;message", subcommand=None, usage="", description="Shows information about a proxied message. Requires the message ID.") async def message_info(conn, message, args): try: @@ -357,14 +376,16 @@ async def message_info(conn, message, args): embed = discord.Embed() embed.timestamp = message.timestamp embed.colour = discord.Colour.blue() - + if system["name"]: system_value = "`{}`: {}".format(system["hid"], system["name"]) else: system_value = "`{}`".format(system["hid"]) embed.add_field(name="System", value=system_value) - embed.add_field(name="Member", value="`{}`: {}".format(member["hid"], member["name"])) - embed.add_field(name="Sent by", value="{}#{}".format(original_sender.name, original_sender.discriminator)) + embed.add_field(name="Member", value="`{}`: {}".format( + member["hid"], member["name"])) + embed.add_field(name="Sent by", value="{}#{}".format( + original_sender.name, original_sender.discriminator)) embed.add_field(name="Content", value=message.clean_content, inline=False) embed.set_author(name=member["name"], url=member["avatar_url"]) @@ -372,12 +393,14 @@ async def message_info(conn, message, args): await client.send_message(message.channel, embed=embed) return True + @command(cmd="pk;help", subcommand=None, usage="[system|member|message]", description="Shows this help message.") async def show_help(conn, message, args): embed = discord.Embed() embed.colour = discord.Colour.blue() embed.title = "PluralKit Help" - embed.set_footer(text="<> denotes mandatory arguments, [] denotes optional arguments") + embed.set_footer( + text="<> denotes mandatory arguments, [] denotes optional arguments") if len(args) > 0 and ("pk;" + args[0]) in command_map: cmds = ["", ("pk;" + args[0], command_map["pk;" + args[0]])] @@ -386,7 +409,8 @@ async def show_help(conn, message, args): for cmd, subcommands in cmds: for subcmd, (_, usage, description) in subcommands.items(): - embed.add_field(name="{} {} {}".format(cmd, subcmd or "", usage or ""), value=description, inline=False) - + embed.add_field(name="{} {} {}".format( + cmd, subcmd or "", usage or ""), value=description, inline=False) + await client.send_message(message.channel, embed=embed) - return True \ No newline at end of file + return True diff --git a/bot/pluralkit/db.py b/bot/pluralkit/db.py index eeda07c2..9244d2d2 100644 --- a/bot/pluralkit/db.py +++ b/bot/pluralkit/db.py @@ -5,6 +5,7 @@ import asyncpg.exceptions from pluralkit.bot import logger + async def connect(): while True: try: @@ -12,32 +13,40 @@ async def connect(): except (ConnectionError, asyncpg.exceptions.CannotConnectNowError): pass + def db_wrap(func): async def inner(*args, **kwargs): before = time.perf_counter() res = await func(*args, **kwargs) after = time.perf_counter() - + logger.debug(" - DB took {:.2f} ms".format((after - before) * 1000)) return res return inner + webhook_cache = {} + @db_wrap async def create_system(conn, system_name: str, system_hid: str): - logger.debug("Creating system (name={}, hid={})".format(system_name, system_hid)) + logger.debug("Creating system (name={}, hid={})".format( + system_name, system_hid)) return await conn.fetchrow("insert into systems (name, hid) values ($1, $2) returning *", system_name, system_hid) + @db_wrap async def remove_system(conn, system_id: int): logger.debug("Deleting system (id={})".format(system_id)) await conn.execute("delete from systems where id = $1", system_id) + @db_wrap async def create_member(conn, system_id: int, member_name: str, member_hid: str): - logger.debug("Creating member (system={}, name={}, hid={})".format(system_id, member_name, member_hid)) - return await conn.fetchrow("insert into members (name, system, hid) values ($1, $2, $3) returning *", member_name, system_id, member_hid) + logger.debug("Creating member (system={}, name={}, hid={})".format( + system_id, member_name, member_hid)) + return await conn.fetchrow("insert into members (name, system, hid) values ($1, $2, $3) returning *", member_name, system_id, member_hid) + @db_wrap async def delete_member(conn, member_id: int): @@ -45,70 +54,90 @@ async def delete_member(conn, member_id: int): await conn.execute("update switches set member = null, member_del = true where member = $1", member_id) await conn.execute("delete from members where id = $1", member_id) + @db_wrap async def link_account(conn, system_id: int, account_id: str): - logger.debug("Linking account (account_id={}, system_id={})".format(account_id, system_id)) + logger.debug("Linking account (account_id={}, system_id={})".format( + account_id, system_id)) await conn.execute("insert into accounts (uid, system) values ($1, $2)", int(account_id), system_id) + @db_wrap async def unlink_account(conn, system_id: int, account_id: str): - logger.debug("Unlinking account (account_id={}, system_id={})".format(account_id, system_id)) + logger.debug("Unlinking account (account_id={}, system_id={})".format( + account_id, system_id)) await conn.execute("delete from accounts where uid = $1 and system = $2", int(account_id), system_id) + @db_wrap async def get_linked_accounts(conn, system_id: int): return [row["uid"] for row in await conn.fetch("select uid from accounts where system = $1", system_id)] + @db_wrap async def get_system_by_account(conn, account_id: str): return await conn.fetchrow("select systems.* from systems, accounts where accounts.uid = $1 and accounts.system = systems.id", int(account_id)) + @db_wrap async def get_system_by_hid(conn, system_hid: str): return await conn.fetchrow("select * from systems where hid = $1", system_hid) + @db_wrap async def get_system(conn, system_id: int): return await conn.fetchrow("select * from systems where id = $1", system_id) + @db_wrap async def get_member_by_name(conn, system_id: int, member_name: str): return await conn.fetchrow("select * from members where system = $1 and name = $2", system_id, member_name) + @db_wrap async def get_member_by_hid_in_system(conn, system_id: int, member_hid: str): return await conn.fetchrow("select * from members where system = $1 and hid = $2", system_id, member_hid) + @db_wrap async def get_member_by_hid(conn, member_hid: str): return await conn.fetchrow("select * from members where hid = $1", member_hid) + @db_wrap async def get_member(conn, member_id: int): return await conn.fetchrow("select * from members where id = $1", member_id) + @db_wrap async def get_message(conn, message_id: str): return await conn.fetchrow("select * from messages where mid = $1", message_id) + @db_wrap async def update_system_field(conn, system_id: int, field: str, value): - logger.debug("Updating system field (id={}, {}={})".format(system_id, field, value)) + logger.debug("Updating system field (id={}, {}={})".format( + system_id, field, value)) await conn.execute("update systems set {} = $1 where id = $2".format(field), value, system_id) + @db_wrap async def update_member_field(conn, member_id: int, field: str, value): - logger.debug("Updating member field (id={}, {}={})".format(member_id, field, value)) + logger.debug("Updating member field (id={}, {}={})".format( + member_id, field, value)) await conn.execute("update members set {} = $1 where id = $2".format(field), value, member_id) + @db_wrap async def get_all_members(conn, system_id: int): return await conn.fetch("select * from members where system = $1", system_id) + @db_wrap async def get_members_exceeding(conn, system_id: int, length: int): return await conn.fetch("select * from members where system = $1 and length(name) >= $2", system_id, length) + @db_wrap async def get_webhook(conn, channel_id: str): if channel_id in webhook_cache: @@ -117,30 +146,38 @@ async def get_webhook(conn, channel_id: str): webhook_cache[channel_id] = res return res + @db_wrap async def add_webhook(conn, channel_id: str, webhook_id: str, webhook_token: str): - logger.debug("Adding new webhook (channel={}, webhook={}, token={})".format(channel_id, webhook_id, webhook_token)) + logger.debug("Adding new webhook (channel={}, webhook={}, token={})".format( + channel_id, webhook_id, webhook_token)) await conn.execute("insert into webhooks (channel, webhook, token) values ($1, $2, $3)", int(channel_id), int(webhook_id), webhook_token) + @db_wrap async def add_message(conn, message_id: str, channel_id: str, member_id: int, sender_id: str): - logger.debug("Adding new message (id={}, channel={}, member={}, sender={})".format(message_id, channel_id, member_id, sender_id)) + logger.debug("Adding new message (id={}, channel={}, member={}, sender={})".format( + message_id, channel_id, member_id, sender_id)) await conn.execute("insert into messages (mid, channel, member, sender) values ($1, $2, $3, $4)", int(message_id), int(channel_id), member_id, int(sender_id)) + @db_wrap async def get_members_by_account(conn, account_id: str): # Returns a "chimera" object return await conn.fetch("select members.id, members.hid, members.prefix, members.suffix, members.name, members.avatar_url, systems.tag from systems, members, accounts where accounts.uid = $1 and systems.id = accounts.system and members.system = systems.id", int(account_id)) + @db_wrap async def get_message_by_sender_and_id(conn, message_id: str, sender_id: str): await conn.fetchrow("select * from messages where mid = $1 and sender = $2", int(message_id), int(sender_id)) + @db_wrap async def delete_message(conn, message_id: str): logger.debug("Deleting message (id={})".format(message_id)) await conn.execute("delete from messages where mid = $1", int(message_id)) + async def create_tables(conn): await conn.execute("""create table if not exists systems ( id serial primary key, @@ -190,4 +227,4 @@ async def create_tables(conn): id bigint primary key, cmd_chans bigint[], proxy_chans bigint[] - )""") \ No newline at end of file + )""") diff --git a/bot/pluralkit/proxy.py b/bot/pluralkit/proxy.py index f99e6b9d..5a87b690 100644 --- a/bot/pluralkit/proxy.py +++ b/bot/pluralkit/proxy.py @@ -6,6 +6,7 @@ import aiohttp from pluralkit import db from pluralkit.bot import client, logger + async def get_webhook(conn, channel): async with conn.transaction(): # Try to find an existing webhook @@ -14,7 +15,8 @@ async def get_webhook(conn, channel): if not hook_row: async with aiohttp.ClientSession() as session: req_data = {"name": "PluralKit Proxy Webhook"} - req_headers = {"Authorization": "Bot {}".format(os.environ["TOKEN"])} + req_headers = { + "Authorization": "Bot {}".format(os.environ["TOKEN"])} async with session.post("https://discordapp.com/api/v6/channels/{}/webhooks".format(channel.id), json=req_data, headers=req_headers) as resp: data = await resp.json() @@ -24,17 +26,19 @@ async def get_webhook(conn, channel): # Insert new hook into DB await db.add_webhook(conn, channel_id=channel.id, webhook_id=hook_id, webhook_token=token) return hook_id, token - + return hook_row["webhook"], hook_row["token"] + async def proxy_message(conn, member, message, inner): - logger.debug("Proxying message '{}' for member {}".format(inner, member["hid"])) + logger.debug("Proxying message '{}' for member {}".format( + inner, member["hid"])) # Delete the original message await client.delete_message(message) # Get the webhook details hook_id, hook_token = await get_webhook(conn, message.channel) - async with aiohttp.ClientSession() as session: + async with aiohttp.ClientSession() as session: req_data = { "username": "{} {}".format(member["name"], member["tag"] or "").strip(), "avatar_url": member["avatar_url"], @@ -49,13 +53,15 @@ async def proxy_message(conn, member, message, inner): # Insert new message details into the DB await db.add_message(conn, message_id=resp_data["id"], channel_id=message.channel.id, member_id=member["id"], sender_id=message.author.id) + async def handle_proxying(conn, message): # Big fat query to find every member associated with this account # Returned member object has a few more keys (system tag, for example) members = await db.get_members_by_account(conn, account_id=message.author.id) # Sort by specificity (members with both prefix and suffix go higher) - members = sorted(members, key=lambda x: int(bool(x["prefix"])) + int(bool(x["suffix"])), reverse=True) + members = sorted(members, key=lambda x: int( + bool(x["prefix"])) + int(bool(x["suffix"])), reverse=True) msg = message.content for member in members: @@ -71,14 +77,14 @@ async def handle_proxying(conn, message): if msg.startswith(prefix) and msg.endswith(suffix): # Extract the actual message contents sans tags if suffix: - inner_message = message.content[len(prefix):-len(suffix)].strip() + inner_message = message.content[len( + prefix):-len(suffix)].strip() else: # Slicing to -0 breaks, don't do that inner_message = message.content[len(prefix):].strip() await proxy_message(conn, member, message, inner_message) break - async def handle_reaction(conn, reaction, user): @@ -90,4 +96,4 @@ async def handle_reaction(conn, reaction, user): if message: # If so, delete the message and remove it from the DB await db.delete_message(conn, message["mid"]) - await client.delete_message(reaction.message) \ No newline at end of file + await client.delete_message(reaction.message) diff --git a/bot/pluralkit/utils.py b/bot/pluralkit/utils.py index b0aa7175..bae1851d 100644 --- a/bot/pluralkit/utils.py +++ b/bot/pluralkit/utils.py @@ -9,9 +9,11 @@ import discord from pluralkit import db from pluralkit.bot import client, logger + def generate_hid() -> str: return "".join(random.choices(string.ascii_lowercase, k=5)) + async def parse_mention(mention: str) -> discord.User: # First try matching mention format match = re.fullmatch("<@!?(\\d+)>", mention) @@ -27,10 +29,11 @@ async def parse_mention(mention: str) -> discord.User: except (ValueError, discord.NotFound): return None + async def get_system_fuzzy(conn, key) -> asyncpg.Record: if isinstance(key, discord.User): return await db.get_system_by_account(conn, account_id=key.id) - + if isinstance(key, str) and len(key) == 5: return await db.get_system_by_hid(conn, system_hid=key) @@ -39,7 +42,8 @@ async def get_system_fuzzy(conn, key) -> asyncpg.Record: if system: return system return None - + + async def get_member_fuzzy(conn, system_id: int, key: str, system_only=True) -> asyncpg.Record: # First search by hid if system_only: @@ -60,6 +64,8 @@ command_map = {} # Command wrapper # Return True for success, return False for failure # Second parameter is the message it'll send. If just False, will print usage + + def command(cmd, subcommand, usage=None, description=None): def wrap(func): async def wrapper(conn, message, args): @@ -70,12 +76,13 @@ def command(cmd, subcommand, usage=None, description=None): success, msg = res, None else: success, msg = res - + if not success and not msg: # Failure, no message, print usage usage_embed = discord.Embed() usage_embed.colour = discord.Colour.blue() - usage_embed.add_field(name="Usage", value=usage, inline=False) + usage_embed.add_field( + name="Usage", value=usage, inline=False) await client.send_message(message.channel, embed=usage_embed) elif not success: @@ -103,6 +110,8 @@ def command(cmd, subcommand, usage=None, description=None): # Member command wrapper # Tries to find member by first argument # If system_only=False, allows members from other systems by hid + + def member_command(cmd, subcommand, usage=None, description=None, system_only=True): def wrap(func): async def wrapper(conn, message, args): @@ -122,11 +131,12 @@ def member_command(cmd, subcommand, usage=None, description=None, system_only=Tr if member is None: return False, "Can't find member \"{}\".".format(args[0]) - + return await func(conn, message, member, args[1:]) return command(cmd=cmd, subcommand=subcommand, usage=usage, description=description)(wrapper) return wrap + async def generate_system_info_card(conn, system: asyncpg.Record) -> discord.Embed: card = discord.Embed() @@ -134,7 +144,8 @@ async def generate_system_info_card(conn, system: asyncpg.Record) -> discord.Emb card.title = system["name"] if system["description"]: - card.add_field(name="Description", value=system["description"], inline=False) + card.add_field(name="Description", + value=system["description"], inline=False) if system["tag"]: card.add_field(name="Tag", value=system["tag"]) @@ -158,11 +169,13 @@ async def generate_system_info_card(conn, system: asyncpg.Record) -> discord.Emb member_texts.append("`{}`: {}".format(member["hid"], member["name"])) if len(member_texts) > 0: - card.add_field(name="Members", value="\n".join(member_texts), inline=False) + card.add_field(name="Members", value="\n".join( + member_texts), inline=False) card.set_footer(text="System ID: {}".format(system["hid"])) return card + async def generate_member_info_card(conn, member: asyncpg.Record) -> discord.Embed: card = discord.Embed() card.set_author(name=member["name"], icon_url=member["avatar_url"]) @@ -171,18 +184,21 @@ async def generate_member_info_card(conn, member: asyncpg.Record) -> discord.Emb card.colour = int(member["color"], 16) if member["birthday"]: - card.add_field(name="Birthdate", value=member["birthday"].strftime("%b %d, %Y")) - + card.add_field(name="Birthdate", + value=member["birthday"].strftime("%b %d, %Y")) + if member["pronouns"]: card.add_field(name="Pronouns", value=member["pronouns"]) if member["prefix"] or member["suffix"]: prefix = member["prefix"] or "" suffix = member["suffix"] or "" - card.add_field(name="Proxy Tags", value="{}text{}".format(prefix, suffix)) + card.add_field(name="Proxy Tags", + value="{}text{}".format(prefix, suffix)) if member["description"]: - card.add_field(name="Description", value=member["description"], inline=False) + card.add_field(name="Description", + value=member["description"], inline=False) # Get system name and hid system = await db.get_system(conn, system_id=member["system"]) @@ -192,9 +208,11 @@ async def generate_member_info_card(conn, member: asyncpg.Record) -> discord.Emb system_value = "`{}`".format(system["hid"]) card.add_field(name="System", value=system_value, inline=False) - card.set_footer(text="System ID: {} | Member ID: {}".format(system["hid"], member["hid"])) + card.set_footer(text="System ID: {} | Member ID: {}".format( + system["hid"], member["hid"])) return card + async def text_input(message, subject): await client.send_message(message.channel, "Reply in this channel with the new description you want to set for {}.".format(subject)) msg = await client.wait_for_message(author=message.author, channel=message.channel) @@ -208,4 +226,4 @@ async def text_input(message, subject): return msg.content else: await client.clear_reactions(msg) - return None \ No newline at end of file + return None