tuxbot-bot/tuxbot/cogs/Polls/polls.py

282 lines
8.6 KiB
Python
Raw Normal View History

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.utils.functions.utils import upper_first
from tuxbot.core.bot import Tux
from tuxbot.core.i18n import (
Translator,
)
from .functions import emotes as utils_emotes, listeners
from .functions.converters import NewPropositionConvertor, PollConverter
from .models import Poll
from .models.suggests import Suggest
log = logging.getLogger("tuxbot.cogs.Polls")
_ = Translator("Polls", __file__)
class Polls(commands.Cog, name="Polls"):
def __init__(self, bot: Tux):
self.bot = bot
async def cog_command_error(self, ctx, error):
await listeners.cog_command_error(self, ctx, error)
@commands.Cog.listener()
async def on_raw_reaction_add(self, pld: discord.RawReactionActionEvent):
await listeners.on_raw_reaction_add(self, pld)
@commands.Cog.listener()
async def on_raw_reaction_remove(
self, pld: discord.RawReactionActionEvent
):
await listeners.on_raw_reaction_remove(self, pld)
# =========================================================================
# =========================================================================
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]} - {upper_first(answer)}__",
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["data"]["labels"].append(
field["name"][6:].replace("__", "")
)
2021-01-21 16:13:46 +01:00
chart_options["data"]["datasets"][0]["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)
async def created_suggest(
self, ctx: ContextPlus, poll: PollConverter, new: str
):
stmt = await ctx.send(
_(
"**Preparation**",
ctx,
self.bot.config,
)
)
suggest_row = await Suggest()
suggest_row.channel_id = ctx.channel.id
suggest_row.message_id = stmt.id
suggest_row.author_id = ctx.author.id
if isinstance(poll, Poll):
# pylint: disable=pointless-string-statement
"""Just to change type for PyCharm"""
suggest_row.poll = poll
suggest_row.proposition = new
await suggest_row.save()
poll_channel: discord.TextChannel = await self.bot.fetch_channel(
poll.channel_id
)
poll_message = await poll_channel.fetch_message(poll.message_id)
e = discord.Embed(
title=_(
"Proposed addition for poll #{id}", ctx, self.bot.config
).format(id=poll.id),
description=new,
)
e.add_field(
name=_("Poll", ctx, self.bot.config),
value="[{}]({})".format(
_("here", ctx, self.bot.config), poll_message.jump_url
),
)
e.set_footer(
text=_("Requested by {author}", ctx, self.bot.config).format(
author=ctx.author.name
),
icon_url=ctx.author.avatar_url_as(format="png"),
)
await stmt.edit(content="", embed=e)
for emote in utils_emotes.check:
await stmt.add_reaction(emote)
async def get_suggest(
self, pld: discord.RawReactionActionEvent
) -> Union[bool, Suggest]:
if pld.user_id != self.bot.user.id:
suggest = await Suggest.get_or_none(message_id=pld.message_id)
if suggest is not None and pld.emoji.name in utils_emotes.check:
return suggest
return False
# =========================================================================
# =========================================================================
@group_extra(name="polls", aliases=["poll", "sondages", "sondage"])
async def _poll(self, ctx: commands.Context):
if ctx.invoked_subcommand is None:
await ctx.send_help("poll")
@_poll.command(name="create", aliases=["new", "nouveau"])
async def _poll_create(self, ctx: ContextPlus, *, poll: str):
args: list = poll.split()
is_anonymous = False
if "--anonymous" in map(str.lower, args):
is_anonymous = True
args.remove("--anonymous")
elif "--anonyme" in map(str.lower, args):
is_anonymous = True
args.remove("--anonyme")
if args[-1] != "|":
args.append("|")
delimiters = [i for i, val in enumerate(args) if val == "|"]
question = upper_first(" ".join(args[: delimiters[0]]))
answers = []
for i in range(len(delimiters) - 1):
start = delimiters[i] + 1
end = delimiters[i + 1]
answers.append(upper_first(" ".join(args[start:end])))
await self.create_poll(ctx, question, answers, anonymous=is_anonymous)
@_poll.command(name="propose", aliases=["suggest", "ajout"])
async def _poll_propose(
self,
ctx: ContextPlus,
poll: PollConverter,
*,
new: NewPropositionConvertor,
):
await self.created_suggest(ctx, poll, str(new))