diff --git a/tuxbot/cogs/Polls/polls.py b/tuxbot/cogs/Polls/polls.py index 21332ab..abfdd1e 100644 --- a/tuxbot/cogs/Polls/polls.py +++ b/tuxbot/cogs/Polls/polls.py @@ -1,12 +1,18 @@ +import json import logging +from typing import Union +import discord from discord.ext import commands +from yarl import URL from tuxbot.core.utils.functions.extra import ContextPlus, group_extra from tuxbot.core.bot import Tux from tuxbot.core.i18n import ( Translator, ) +from .functions import emotes as utils_emotes +from .models import Poll, Response log = logging.getLogger("tuxbot.cogs.Polls") _ = Translator("Polls", __file__) @@ -16,36 +22,216 @@ class Polls(commands.Cog, name="Polls"): def __init__(self, bot: Tux): self.bot = bot + @commands.Cog.listener() + async def on_raw_reaction_add(self, pld: discord.RawReactionActionEvent): + poll = await self.get_poll(pld) + + if poll: + if poll.is_anonymous: + try: + await self.remove_reaction(pld) + except discord.errors.Forbidden: + pass + choice = utils_emotes.get_index(pld.emoji.name) + + response = await Response.get_or_none( + user_id=pld.user_id, choice=choice, poll__id=poll.id + ) + + if response is not None: + await poll.choices.remove(response) + await response.delete() + else: + res = await Response.create( + user_id=pld.user_id, poll=poll, choice=choice + ) + await poll.choices.add(res) + + await self.update_poll(poll) + + @commands.Cog.listener() + async def on_raw_reaction_remove( + self, pld: discord.RawReactionActionEvent + ): + poll = await self.get_poll(pld) + + if poll: + choice = utils_emotes.get_index(pld.emoji.name) + + response = await Response.get_or_none( + user_id=pld.user_id, choice=choice, poll__id=poll.id + ) + + if response is not None: + await poll.choices.remove(response) + await response.delete() + + await self.update_poll(poll) + + # ========================================================================= + # ========================================================================= + + async def create_poll( + self, + ctx: ContextPlus, + question: str, + answers: list[str], + anonymous=False, + ): + emotes = utils_emotes.get(len(answers)) + + stmt = await ctx.send( + _( + "**Preparation**", + ctx, + self.bot.config, + ) + ) + + poll_row = await Poll() + + poll_row.channel_id = stmt.channel.id + poll_row.message_id = stmt.id + poll_row.author_id = ctx.author.id + poll_row.content = {} + poll_row.is_anonymous = anonymous + poll_row.available_choices = len(answers) + + await poll_row.save() + + e = discord.Embed(description=f"**{question}**") + e.set_author( + name=ctx.author, + icon_url="https://img.icons8.com/plasticine/100/000000/survey.png", + ) + for i, answer in enumerate(answers): + e.add_field( + name=f"__{emotes[i]} - {answer.capitalize()}__", + value="**0** vote", + ) + e.set_footer(text=f"ID: #{poll_row.id}") + + poll_row.content = e.to_dict() + await poll_row.save() + + await stmt.edit(content="", embed=e) + for emote in range(len(answers)): + await stmt.add_reaction(emotes[emote]) + + async def get_poll( + self, pld: discord.RawReactionActionEvent + ) -> Union[bool, Poll]: + if pld.user_id != self.bot.user.id: + poll = await Poll.get_or_none(message_id=pld.message_id) + + if poll is not None: + emotes = utils_emotes.get(poll.available_choices) + + if pld.emoji.name in emotes: + return poll + + return False + + async def update_poll(self, poll: Poll): + channel: discord.TextChannel = self.bot.get_channel(poll.channel_id) + message: discord.Message = await channel.fetch_message(poll.message_id) + + chart_base_url = "https://quickchart.io/chart?backgroundColor=white&c=" + chart_options = { + "type": "pie", + "data": {"labels": [], "datasets": [{"data": []}]}, + } + + content = ( + json.loads(poll.content) + if isinstance(poll.content, str) + else poll.content + ) + + responses = {} + + async for response in poll.choices: + if responses.get(response.choice): + responses[response.choice] += 1 + else: + responses[response.choice] = 1 + + for i, field in enumerate(content.get("fields")): + responders = responses.get(i, 0) + + chart_options.get("data").get("labels").append( + field.get("name")[6:].replace("__", "") + ) + + chart_options.get("data").get("datasets")[0].get("data").append( + responders + ) + + if responders <= 1: + field["value"] = f"**{responders}** vote" + else: + field["value"] = f"**{responders}** votes" + + e = discord.Embed(description=content.get("description")) + e.set_author( + name=content.get("author").get("name"), + icon_url=content.get("author").get("icon_url"), + ) + chart_url = URL(chart_base_url + json.dumps(chart_options)) + e.set_thumbnail(url=str(chart_url)) + + for field in content.get("fields"): + e.add_field( + name=field.get("name"), value=field.get("value"), inline=True + ) + + e.set_footer(text=content.get("footer").get("text")) + + await message.edit(embed=e) + + poll.content = json.dumps(content) + + await poll.save() + + async def remove_reaction(self, pld: discord.RawReactionActionEvent): + channel: discord.TextChannel = self.bot.get_channel(pld.channel_id) + message: discord.Message = await channel.fetch_message(pld.message_id) + user: discord.User = await self.bot.fetch_user(pld.user_id) + + await message.remove_reaction(pld.emoji.name, user) + # ========================================================================= # ========================================================================= @group_extra(name="polls", aliases=["poll", "sondages", "sondage"]) - async def _polls(self, ctx: ContextPlus, *, message): + async def _poll(self, ctx: commands.Context): if ctx.invoked_subcommand is None: - args: list = message.lower().split() - is_anonymous = False + await ctx.send_help("poll") - if "--anonymous" in args: - is_anonymous = True - args.remove("--anonymous") - elif "--anonyme" in args: - is_anonymous = True - args.remove("--anonyme") + @_poll.group(name="create", aliases=["new", "nouveau"]) + async def _poll_create(self, ctx: ContextPlus, *, poll: str): + args: list = poll.lower().split() + is_anonymous = False - if args[-1] != "|": - args.append("|") + if "--anonymous" in args: + is_anonymous = True + args.remove("--anonymous") + elif "--anonyme" in args: + is_anonymous = True + args.remove("--anonyme") - delimiters = [i for i, val in enumerate(args) if val == "|"] + if args[-1] != "|": + args.append("|") - question = " ".join(args[0 : delimiters[0]]).capitalize() - answers = [] + delimiters = [i for i, val in enumerate(args) if val == "|"] - for i in range(len(delimiters) - 1): - start = delimiters[i] + 1 - end = delimiters[i + 1] + question = " ".join(args[: delimiters[0]]).capitalize() + answers = [] - answers.append(" ".join(args[start:end]).capitalize()) + for i in range(len(delimiters) - 1): + start = delimiters[i] + 1 + end = delimiters[i + 1] - await ctx.send( - f"{message=}\n{question=}\n{answers=}\n{is_anonymous=}" - ) + answers.append(" ".join(args[start:end]).capitalize()) + + await self.create_poll(ctx, question, answers, anonymous=is_anonymous)