From 43a26de389abfb7da469942acfa1cef0dc55af46 Mon Sep 17 00:00:00 2001 From: konsumlamm Date: Mon, 28 Sep 2020 17:43:30 +0200 Subject: Refactor project to use the command extension Restart bot when an exception occurs Add .idea/ to .gitignore --- .gitignore | 1 + bot.py | 167 +++++++++++++------------------------------------------ command_utils.py | 62 --------------------- commands.py | 107 +++++++++++++++++++++++++++++++++++ config.py | 2 +- 5 files changed, 148 insertions(+), 191 deletions(-) delete mode 100644 command_utils.py create mode 100644 commands.py diff --git a/.gitignore b/.gitignore index 4c49bd7..a9ad188 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ .env +.idea/ diff --git a/bot.py b/bot.py index 88153cc..75f8aea 100644 --- a/bot.py +++ b/bot.py @@ -1,19 +1,13 @@ -import asyncio -import discord -import config -from command_utils import command, CommandClient -from random import shuffle, choice -from time import sleep +from os import getenv +from commands import answer, await_n, bot_commands +import config -async def await_n(lst): - lst = list(asyncio.create_task(task) for task in lst) - if lst: - done, _ = await asyncio.wait(lst) - return list(task.result() for task in done) +import discord +from discord.ext import commands -class Client(CommandClient): +class Cupido(commands.Bot): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.meta_channel = None @@ -23,45 +17,38 @@ class Client(CommandClient): async def on_ready(self): print(f'the bot {config.NAME} is logged in as "{self.user}"') - @command( - names = ('help', 'hepl', 'h', '?'), - description = 'display this help message', - is_help = True - ) - async def help(self, ctx): - command_doc = '\n'.join( - f' * {config.COMMAND_PREFIX.strip()} {c.names[0]:15} - {c.description}' - for c in self.get_commands()) - await ctx.answer(f'''``` -{config.HELP_TEXT}\nThese are all available commands:\n{command_doc}```''') - async def create_category(self, ctx): category = await ctx.guild.create_category(config.CATEGORY_CHANNEL_NAME) - await ctx.answer(f'info: category created "{category}" ({category.id})') + await answer(ctx, f'info: category created "{category}" ({category.id})') return category async def create_lobby(self, ctx): - lobby = await ctx.guild.create_voice_channel(config.LOBBY_CHANNEL_NAME, topic=config.LOBBY_CHANNEL_TOPIC, category=self.meta_channel) - await ctx.answer(f'info: voice channel created "{lobby}" ({lobby.id})') + lobby = await ctx.guild.create_voice_channel( + config.LOBBY_CHANNEL_NAME, + topic=config.LOBBY_CHANNEL_TOPIC, + category=self.meta_channel, + ) + await answer(ctx, f'info: voice channel created "{lobby}" ({lobby.id})') return lobby def get_meta_channel(self, ctx): - return (self.meta_channel - or discord.utils.get(ctx.guild.categories, - name=config.CATEGORY_CHANNEL_NAME)) + return self.meta_channel or discord.utils.get(ctx.guild.categories, name=config.CATEGORY_CHANNEL_NAME) def get_lobby_channel(self, ctx, meta_channel): return (self.lobby_channel - or discord.utils.get(ctx.guild.voice_channels, - name=config.LOBBY_CHANNEL_NAME, - category=meta_channel)) + or discord.utils.get( + ctx.guild.voice_channels, + name=config.LOBBY_CHANNEL_NAME, + category=meta_channel + )) def get_pair_channels(self, ctx, meta_channel): return (self.pair_channels or sorted((channel for channel in ctx.guild.voice_channels if channel.category == meta_channel and channel.name.isdigit()), - key=lambda c: c.name)) + key=lambda c: c.name, + )) async def destroy_pair_channels(self, ctx, meta_channel): await await_n(channel.delete() for channel in self.get_pair_channels(ctx, meta_channel)) @@ -74,110 +61,34 @@ class Client(CommandClient): futures.append(ctx.guild.create_voice_channel(str(i), category=meta_channel)) return await await_n(futures) - @command( - names = ('init', 'create', 'inti', 'craete', 'cretae', 'c', 'i', '+'), - description = 'create a new lobby' - ) - async def init(self, ctx): - self.meta_channel = ( - self.get_meta_channel(ctx) - or await self.create_category(ctx) - ) - self.lobby_channel = ( - self.get_lobby_channel(ctx, self.meta_channel) - or await self.create_lobby(ctx) - ) - - @command( - names = ('destroy', 'kill', 'desctruction', 'genocide', '-'), - description = f'destruct all {config.NAME} channels' - ) - async def destroy(self, ctx): - futures = [] - meta_channel = self.get_meta_channel(ctx) - for channel in (self.get_lobby_channel(ctx, meta_channel), meta_channel): - if channel: - futures.append(channel.delete()) - await await_n(futures) - self.lobby_channel = None - self.meta_channel = None - self.pair_channels = [] - async def get_channels(self, ctx): meta_channel = self.get_meta_channel(ctx) lobby_channel = self.get_lobby_channel(ctx, meta_channel) if meta_channel is None or lobby_channel is None: - await ctx.answer('error: cannot start shuffling, you need to initialize channels') - await self.help(ctx) + await answer(ctx, 'error: cannot start shuffling, you need to initialize channels') return None return meta_channel, lobby_channel - @command( - names = ('shuffle', 'start', 'run', 'strat', 'rnu'), - description = 'start shuffling' - ) - async def shuffle(self, ctx): - channels = await self.get_channels(ctx) - if not channels: return - meta_channel, lobby_channel = channels - members = lobby_channel.members[:] - slots = len(members) >> 1 - self.pair_channels = await self.create_pair_channels(ctx, meta_channel, slots) - slots = [] - for i, _ in enumerate(self.pair_channels): - slots.append(i) - slots.append(i) - shuffle(slots) - futures = [] - for slot in slots: - member = members.pop() - if member is None: break - futures.append(member.move_to(self.pair_channels[slot])) - if members: - futures.append(members.pop().move_to(choice(self.pair_channels))) - await await_n(futures) - - @command( - names = ('stop', 'quit', 'exit', 'abort', 'back', 'return'), - description = 'move everyone back to lobby' - ) - async def stop(self, ctx): - channels = await self.get_channels(ctx) - if not channels: return - meta_channel, lobby_channel = channels - pair_channels = self.get_pair_channels(ctx, meta_channel) - futures = [] - for channel in pair_channels: - for member in channel.members: - futures.append(member.move_to(lobby_channel)) - await await_n(futures) - await self.destroy_pair_channels(ctx, meta_channel) - - @command( - names = ('loop',), - description = 'repeat "shuffle" and "stop" (default: 3) times and (default: 120) seconds' - ) - async def loop(self, ctx): - if len(ctx.args) >= 1 and ctx.args[0].isdigit(): - n = int(ctx.args[0]) - else: - n = 3 - if len(ctx.args) >= 2 and ctx.args[1].isdigit(): - t = int(ctx.args[1]) - else: - t = 120 - await ctx.answer(f'repeat shuffling {n} times and each {t} seconds') - for _ in range(n): - await self.shuffle(ctx) - sleep(t) - await self.stop(ctx) - -if __name__ == '__main__': - from os import getenv +def main(): token = getenv(config.TOKEN_ENV_VAR) if token is None: print('error: no token was given') exit(1) - bot = Client(activity=discord.Game(name=config.GAME_STATUS)) + bot = Cupido(activity=discord.Game(name=config.GAME_STATUS), command_prefix=config.COMMAND_PREFIX) + bot.remove_command('help') + for cmd in bot_commands: + bot.add_command(cmd) bot.run(token) + + +if __name__ == '__main__': + while True: + try: + main() + except Exception as e: + print(f'[DEBUG] encountered exception: {type(e)}: {e}') + if e.args == ('Event loop is closed',): + print('[DEBUG] quit...') + break + print('[DEBUG] restarting...') diff --git a/command_utils.py b/command_utils.py deleted file mode 100644 index dc7f433..0000000 --- a/command_utils.py +++ /dev/null @@ -1,62 +0,0 @@ -import discord - -import config - - -def istrue(value, key): - return hasattr(value, key) and getattr(value, key) - - -class Context: - def __init__(self, msg, args): - self.msg = msg - self.args = args - self.guild = msg.guild - - async def answer(self, value): - await self.msg.channel.send(f'{self.msg.author.mention} {value}') - - -class CommandClientMeta(type): - def __new__(cls, name, bases, dct): - commands = [] - for item in dct.copy().values(): - if callable(item) and istrue(item, 'command'): - if istrue(item, 'help_command'): - dct['help_command'] = item - commands.append(item) - dct['commands'] = commands - return super().__new__(cls, name, bases, dct) - - -class CommandClient(discord.Client, metaclass=CommandClientMeta): - def get_commands(self): - return type(self).commands.copy() - - async def run_help(self, ctx): - return await type(self).help_command(self, ctx) - - async def on_message(self, msg): - text = msg.content.strip() - if not text.startswith(config.COMMAND_PREFIX): - return - try: - cmd, *args = [v for v in text[len(config.COMMAND_PREFIX):].split(' ') if v] - except ValueError: - return await self.run_help(Context(msg, [])) - ctx = Context(msg, args) - for command in self.get_commands(): - if cmd in command.names: - return await command(self, ctx) - return await self.run_help(ctx) - - -def command(names=[], description='', is_help=False): - def meta(f): - f.command = True - f.names = names - f.description = description - if is_help: - f.help_command = is_help - return f - return meta diff --git a/commands.py b/commands.py new file mode 100644 index 0000000..e463b98 --- /dev/null +++ b/commands.py @@ -0,0 +1,107 @@ +import asyncio +from typing import Any, List +import random +from time import sleep + +import config + +from discord.ext import commands + + +async def await_n(it) -> List[Any]: + lst = [asyncio.create_task(task) for task in it] + if not lst: + return [] + done, _ = await asyncio.wait(lst) + return [task.result() for task in done] + + +async def answer(ctx: commands.Context, msg: str): + await ctx.send(f'{ctx.author.mention} {msg}') + + +@commands.command(help='display this help message', aliases=('hepl', 'h', '?')) +async def help(ctx: commands.Context): + command_doc = '\n'.join( + f' * {config.COMMAND_PREFIX.strip()} {c.name:15} - {c.help}' + for c in ctx.bot.commands) + await answer(ctx, f'''``` +{config.HELP_TEXT}\nThese are all available commands:\n{command_doc}```''') + + +@commands.command(help='create a new lobby', aliases=('create', 'inti', 'craete', 'cretae', 'c', 'i', '+')) +async def init(ctx: commands.Context): + ctx.bot.meta_channel = ctx.bot.get_meta_channel(ctx) or await ctx.bot.create_category(ctx) + ctx.bot.lobby_channel = ctx.bot.get_lobby_channel(ctx, ctx.bot.meta_channel) or await ctx.bot.create_lobby(ctx) + + +@commands.command(help=f'destruct all {config.NAME} channels', aliases=('kill', 'desctruction', 'genocide', '-')) +async def destroy(ctx: commands.Context): + futures = [] + meta_channel = ctx.bot.get_meta_channel(ctx) + for channel in (ctx.bot.get_lobby_channel(ctx, meta_channel), meta_channel): + if channel: + futures.append(channel.delete()) + await await_n(futures) + ctx.bot.lobby_channel = None + ctx.bot.meta_channel = None + ctx.bot.pair_channels = [] + + +@commands.command(help='start shuffling', aliases=('start', 'run', 'strat', 'rnu')) +async def shuffle(ctx: commands.Context): + channels = await ctx.bot.get_channels(ctx) + if not channels: + return + meta_channel, lobby_channel = channels + members = lobby_channel.members[:] + slots = len(members) >> 1 + ctx.bot.pair_channels = await ctx.bot.create_pair_channels(ctx, meta_channel, slots) + slots = [] + for i, _ in enumerate(ctx.bot.pair_channels): + slots.append(i) + slots.append(i) + random.shuffle(slots) + futures = [] + for slot in slots: + member = members.pop() + if member is None: break + futures.append(member.move_to(ctx.bot.pair_channels[slot])) + if members: + futures.append(members.pop().move_to(random.choice(ctx.bot.pair_channels))) + await await_n(futures) + + +@commands.command(help='move everyone back to lobby', aliases=('quit', 'exit', 'abort', 'back', 'return')) +async def stop(ctx: commands.Context): + channels = await ctx.bot.get_channels(ctx) + if not channels: + return + meta_channel, lobby_channel = channels + pair_channels = ctx.bot.get_pair_channels(ctx, meta_channel) + futures = [] + for channel in pair_channels: + for member in channel.members: + futures.append(member.move_to(lobby_channel)) + await await_n(futures) + await ctx.bot.destroy_pair_channels(ctx, meta_channel) + + +@commands.command(help='repeat "shuffle" and "stop" (default: 3) times and (default: 120) seconds') +async def loop(ctx: commands.Context, *args): + if len(args) >= 1 and args[0].isdigit(): + n = int(args[0]) + else: + n = 3 + if len(args) >= 2 and args[1].isdigit(): + t = int(ctx.args[1]) + else: + t = 120 + await answer(ctx, f'repeat shuffling {n} times and each {t} seconds') + for _ in range(n): + await shuffle(ctx) + sleep(t) + await stop(ctx) + + +bot_commands = [cmd for cmd in locals().values() if isinstance(cmd, commands.Command)] diff --git a/config.py b/config.py index 8ba8767..303b5bf 100644 --- a/config.py +++ b/config.py @@ -1,7 +1,7 @@ NAME = 'cupido' TOKEN_ENV_VAR = 'DISCORD_TOKEN' -COMMAND_PREFIX = '!<3' +COMMAND_PREFIX = '!<3 ' GAME_STATUS = 'with love' -- cgit v1.2.3-54-g00ecf