From 9d54a490c54e10d8f122a6f4bdd14f8af4b4e7c5 Mon Sep 17 00:00:00 2001 From: NatrixAeria Date: Mon, 28 Sep 2020 06:52:22 +0200 Subject: Add shuffle command --- bot.py | 121 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++------ config.py | 2 ++ 2 files changed, 113 insertions(+), 10 deletions(-) diff --git a/bot.py b/bot.py index f97f218..88153cc 100644 --- a/bot.py +++ b/bot.py @@ -2,6 +2,15 @@ import asyncio import discord import config from command_utils import command, CommandClient +from random import shuffle, choice +from time import sleep + + +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) class Client(CommandClient): @@ -9,6 +18,7 @@ class Client(CommandClient): super().__init__(*args, **kwargs) self.meta_channel = None self.lobby_channel = None + self.pair_channels = [] async def on_ready(self): print(f'the bot {config.NAME} is logged in as "{self.user}"') @@ -37,14 +47,32 @@ class Client(CommandClient): def get_meta_channel(self, ctx): return (self.meta_channel - or discord.utils.get(ctx.guild.categories, - name=config.CATEGORY_CHANNEL_NAME)) + or discord.utils.get(ctx.guild.categories, + name=config.CATEGORY_CHANNEL_NAME)) - def get_lobby_channel(self, ctx): + 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=self.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)) + + async def destroy_pair_channels(self, ctx, meta_channel): + await await_n(channel.delete() for channel in self.get_pair_channels(ctx, meta_channel)) + self.pair_channels = [] + + async def create_pair_channels(self, ctx, meta_channel, n): + await self.destroy_pair_channels(ctx, meta_channel) + futures = [] + for i in range(1, n + 1): + 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', '+'), @@ -56,7 +84,7 @@ class Client(CommandClient): or await self.create_category(ctx) ) self.lobby_channel = ( - self.get_lobby_channel(ctx) + self.get_lobby_channel(ctx, self.meta_channel) or await self.create_lobby(ctx) ) @@ -66,10 +94,83 @@ class Client(CommandClient): ) async def destroy(self, ctx): futures = [] - for channel in (self.get_meta_channel(ctx), self.get_lobby_channel(ctx)): + 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 asyncio.wait(futures) + 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) + 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__': @@ -78,5 +179,5 @@ if __name__ == '__main__': if token is None: print('error: no token was given') exit(1) - bot = Client(activity=discord.Game(name=config)) + bot = Client(activity=discord.Game(name=config.GAME_STATUS)) bot.run(token) diff --git a/config.py b/config.py index 404c66a..8ba8767 100644 --- a/config.py +++ b/config.py @@ -3,6 +3,8 @@ TOKEN_ENV_VAR = 'DISCORD_TOKEN' COMMAND_PREFIX = '!<3' +GAME_STATUS = 'with love' + HELP_TEXT = f'''{NAME.title()} - your partner for getting shuffled. {NAME.title()} is a discord bot to get to know your people. This software is open-source ! -- cgit v1.2.3-54-g00ecf