Make loop command interruptable

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" <n> (default: 3) times and <t> (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" <n> (default: {config.DEFAULT_LOOP_COUNT}) times and <t>'
+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 <n> and <t> (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 <https://git.kobert.dev/lovefinderrz.git/>!
 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