From 8883d72fedbb438b0fa14b1dae0405d6317f7e07 Mon Sep 17 00:00:00 2001 From: NatrixAeria Date: Tue, 29 Sep 2020 18:33:21 +0200 Subject: Make loop command interruptable --- bot.py | 43 ++++++++++++++++++++++++++++++++++++------- commands.py | 54 ++++++++++++++++++++++++++++++++++++++---------------- config.py | 5 +++++ 3 files changed, 79 insertions(+), 23 deletions(-) diff --git a/bot.py b/bot.py index 75f8aea..6d8ad9a 100644 --- a/bot.py +++ b/bot.py @@ -1,5 +1,6 @@ from os import getenv +import asyncio from commands import answer, await_n, bot_commands import config @@ -13,6 +14,24 @@ class Cupido(commands.Bot): self.meta_channel = None self.lobby_channel = None self.pair_channels = [] + self.task = None + + async def await_coroutine(self, co): + if self.task is not None: + return 'already running' + self.task = asyncio.create_task(co) + try: + result = await self.task + except asyncio.CancelledError: + self.task = None + return 'cancelled' + finally: + self.task = None + return result + + async def cancel_coroutine(self): + if self.task is not None: + self.task.cancel() async def on_ready(self): print(f'the bot {config.NAME} is logged in as "{self.user}"') @@ -42,16 +61,24 @@ class Cupido(commands.Bot): category=meta_channel )) + def get_pair_channels_no_cache(self, ctx, meta_channel): + return sorted((channel for channel in ctx.guild.voice_channels + if channel.category == meta_channel + and channel.name.isdigit()), + key=lambda c: c.name) + 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, - )) + return self.pair_channels or self.get_pair_channels_no_cache(ctx, meta_channel) + + async def try_delete_channel(self, channel): + try: + await channel.delete() + return True + except discord.errors.NotFound: + return False async def destroy_pair_channels(self, ctx, meta_channel): - await await_n(channel.delete() for channel in self.get_pair_channels(ctx, meta_channel)) + await await_n(map(self.try_delete_channel, self.get_pair_channels_no_cache(ctx, meta_channel))) self.pair_channels = [] async def create_pair_channels(self, ctx, meta_channel, n): @@ -86,6 +113,8 @@ if __name__ == '__main__': while True: try: main() + except KeyboardInterrupt: + break except Exception as e: print(f'[DEBUG] encountered exception: {type(e)}: {e}') if e.args == ('Event loop is closed',): diff --git a/commands.py b/commands.py index e463b98..4e34bce 100644 --- a/commands.py +++ b/commands.py @@ -1,7 +1,6 @@ import asyncio from typing import Any, List import random -from time import sleep import config @@ -52,9 +51,15 @@ async def destroy(ctx: commands.Context): async def shuffle(ctx: commands.Context): channels = await ctx.bot.get_channels(ctx) if not channels: - return + return False meta_channel, lobby_channel = channels members = lobby_channel.members[:] + if len(members) == 0: + await answer(ctx, 'nobody wants to get shuffeled :(') + return False + elif len(members) == 1: + await answer(ctx, f'{members[0].mention} you are a looser and you don\'t have any friends, get a life') + return False slots = len(members) >> 1 ctx.bot.pair_channels = await ctx.bot.create_pair_channels(ctx, meta_channel, slots) slots = [] @@ -70,15 +75,18 @@ async def shuffle(ctx: commands.Context): if members: futures.append(members.pop().move_to(random.choice(ctx.bot.pair_channels))) await await_n(futures) + return True @commands.command(help='move everyone back to lobby', aliases=('quit', 'exit', 'abort', 'back', 'return')) -async def stop(ctx: commands.Context): +async def stop(ctx: commands.Context, cancel_loop=True): 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) + if cancel_loop: + await ctx.bot.cancel_coroutine() futures = [] for channel in pair_channels: for member in channel.members: @@ -87,21 +95,35 @@ async def stop(ctx: commands.Context): 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 +async def loop_cycle(ctx, t): + if not await shuffle(ctx): + return False + await asyncio.sleep(t) + await stop(ctx, cancel_loop=False) + return True + + +@commands.command( + help=f'repeat "shuffle" and "stop" (default: {config.DEFAULT_LOOP_COUNT}) times and ' +f' (default: {config.DEFAULT_LOOP_TIME}) seconds', + aliases=('repeat', 'shuffleloop')) +async def loop(ctx: commands.Context, n=config.DEFAULT_LOOP_COUNT, t=config.DEFAULT_LOOP_TIME): + try: + n, t = int(n), int(t) + except ValueError: + await answer(ctx, f'expecting positive integer arguments for and (e.g. "loop 5 60" to repeat 5 times)') + return await answer(ctx, f'repeat shuffling {n} times and each {t} seconds') for _ in range(n): - await shuffle(ctx) - sleep(t) - await stop(ctx) + result = await ctx.bot.await_coroutine(loop_cycle(ctx, t)) + message = { + 'cancelled': 'cancelled loop command', + 'already running': 'cannot start loop, another task is running...' + }.get(result, None) + if message is not None: + return await answer(ctx, message) + if not result: + break bot_commands = [cmd for cmd in locals().values() if isinstance(cmd, commands.Command)] diff --git a/config.py b/config.py index 303b5bf..0b7987a 100644 --- a/config.py +++ b/config.py @@ -13,3 +13,8 @@ This software is open-source ! CATEGORY_CHANNEL_NAME = NAME.title() LOBBY_CHANNEL_NAME = 'lobby' LOBBY_CHANNEL_TOPIC = "You wanna get shuffled? You're at the right place!" + +# time that one loop cycle needs (in seconds) +DEFAULT_LOOP_TIME = 120 +# time that cycles, that "loop" passes +DEFAULT_LOOP_COUNT = 3 -- cgit v1.2.3-54-g00ecf