397 lines
12 KiB
Python
397 lines
12 KiB
Python
|
from datetime import datetime, timedelta
|
|||
|
|
import logging
|
|||
|
|
import json
|
|||
|
|
from typing import Dict, List, Tuple
|
|||
|
|
from naff import (
|
|||
|
|
Client,
|
|||
|
|
Extension,
|
|||
|
|
slash_command,
|
|||
|
|
slash_option,
|
|||
|
|
InteractionContext,
|
|||
|
|
OptionTypes,
|
|||
|
|
Permissions,
|
|||
|
|
Modal,
|
|||
|
|
ShortText,
|
|||
|
|
ParagraphText,
|
|||
|
|
ModalContext,
|
|||
|
|
Button,
|
|||
|
|
ButtonStyles,
|
|||
|
|
Embed,
|
|||
|
|
spread_to_rows,
|
|||
|
|
Message,
|
|||
|
|
listen,
|
|||
|
|
Task,
|
|||
|
|
IntervalTrigger,
|
|||
|
|
GuildText,
|
|||
|
|
)
|
|||
|
|
from naff.api import events
|
|||
|
|
|
|||
|
|
from database import Polls as PollsModel, PollVotes as PollVotesModel
|
|||
|
|
from peewee import fn
|
|||
|
|
|
|||
|
|
PollOptions = List[Tuple[str | None, str]]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def datetime_to_discord_time(dt: datetime) -> str:
|
|||
|
|
t = dt.strftime("%s")
|
|||
|
|
return f"<t:{int(t)}:R>"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def generate_bar(num: int, total: int, length: int = 10) -> str:
|
|||
|
|
full_char = "\u2593"
|
|||
|
|
empty_char = "\u2591"
|
|||
|
|
|
|||
|
|
if total <= 0:
|
|||
|
|
return length * empty_char
|
|||
|
|
|
|||
|
|
result = round(num / total * length)
|
|||
|
|
percent = num / total * 100
|
|||
|
|
return (
|
|||
|
|
result * full_char
|
|||
|
|
+ (length - result) * empty_char
|
|||
|
|
+ f" ({percent: 3.1f}%)".replace(".0", "")
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def generate_poll_embed(
|
|||
|
|
title: str,
|
|||
|
|
options: PollOptions,
|
|||
|
|
votes: List[int],
|
|||
|
|
*,
|
|||
|
|
multiple_choice: bool = False,
|
|||
|
|
expires: datetime = None,
|
|||
|
|
) -> Embed:
|
|||
|
|
data = []
|
|||
|
|
if multiple_choice:
|
|||
|
|
data.append("\u2022 Multiple choice")
|
|||
|
|
if expires:
|
|||
|
|
data.append(f"\u2022 Expiry: {datetime_to_discord_time(expires)}")
|
|||
|
|
|
|||
|
|
embed = Embed(
|
|||
|
|
title=title,
|
|||
|
|
description=("\n".join(data) if data else None),
|
|||
|
|
)
|
|||
|
|
sum_votes = sum(votes)
|
|||
|
|
|
|||
|
|
for i, (emoji, option) in enumerate(options):
|
|||
|
|
embed.add_field(
|
|||
|
|
name=f"**{emoji if emoji else num_to_emoji(i+1)} {option}**",
|
|||
|
|
value=generate_bar(votes[i], sum_votes) + "\n" + f"{votes[i]} votes",
|
|||
|
|
inline=False,
|
|||
|
|
)
|
|||
|
|
return embed
|
|||
|
|
|
|||
|
|
|
|||
|
|
def num_to_emoji(num: int) -> str:
|
|||
|
|
match num:
|
|||
|
|
case 0:
|
|||
|
|
return "0️⃣"
|
|||
|
|
case 1:
|
|||
|
|
return "1️⃣"
|
|||
|
|
case 2:
|
|||
|
|
return "2️⃣"
|
|||
|
|
case 3:
|
|||
|
|
return "3️⃣"
|
|||
|
|
case 4:
|
|||
|
|
return "4️⃣"
|
|||
|
|
case 5:
|
|||
|
|
return "5️⃣"
|
|||
|
|
case 6:
|
|||
|
|
return "6️⃣"
|
|||
|
|
case 7:
|
|||
|
|
return "7️⃣"
|
|||
|
|
case 8:
|
|||
|
|
return "8️⃣"
|
|||
|
|
case 9:
|
|||
|
|
return "9️⃣"
|
|||
|
|
case 10:
|
|||
|
|
return "🔟"
|
|||
|
|
|
|||
|
|
case _:
|
|||
|
|
raise ValueError(f"Invalid number: `num` must be 0 <= num <= 10.")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Polls(Extension):
|
|||
|
|
def __init__(self, client: Client):
|
|||
|
|
self.client = client
|
|||
|
|
|
|||
|
|
@listen(events.Ready)
|
|||
|
|
async def on_ready(self):
|
|||
|
|
await self.poll_expiry_check()
|
|||
|
|
self.poll_expiry_check.start()
|
|||
|
|
|
|||
|
|
@slash_command(
|
|||
|
|
name="polls",
|
|||
|
|
description="Polls",
|
|||
|
|
sub_cmd_name="create",
|
|||
|
|
sub_cmd_description="Create a poll",
|
|||
|
|
dm_permission=False,
|
|||
|
|
default_member_permissions=Permissions.SEND_MESSAGES,
|
|||
|
|
)
|
|||
|
|
@slash_option(
|
|||
|
|
name="title",
|
|||
|
|
description="Title of the poll",
|
|||
|
|
required=True,
|
|||
|
|
opt_type=OptionTypes.STRING,
|
|||
|
|
)
|
|||
|
|
@slash_option(
|
|||
|
|
name="duration",
|
|||
|
|
description="Duration of the poll in minutes",
|
|||
|
|
required=False,
|
|||
|
|
opt_type=OptionTypes.INTEGER,
|
|||
|
|
)
|
|||
|
|
@slash_option(
|
|||
|
|
name="multiple-choice",
|
|||
|
|
description="If users can vote for multiple options",
|
|||
|
|
required=False,
|
|||
|
|
opt_type=OptionTypes.BOOLEAN,
|
|||
|
|
)
|
|||
|
|
async def create_poll(
|
|||
|
|
self,
|
|||
|
|
ctx: InteractionContext,
|
|||
|
|
*,
|
|||
|
|
title: str,
|
|||
|
|
duration: int | None = None,
|
|||
|
|
multiple_choice: bool = False,
|
|||
|
|
):
|
|||
|
|
modal = Modal(
|
|||
|
|
title="Creating poll",
|
|||
|
|
components=[
|
|||
|
|
ShortText(
|
|||
|
|
custom_id="title",
|
|||
|
|
label="Title",
|
|||
|
|
value=title,
|
|||
|
|
required=True,
|
|||
|
|
max_length=120,
|
|||
|
|
),
|
|||
|
|
ParagraphText(
|
|||
|
|
custom_id="options",
|
|||
|
|
label="Poll options",
|
|||
|
|
placeholder=(
|
|||
|
|
"Add poll options here.\n\n" "- ✅: Yes\n" "- ❌: No\n" "- Unsure"
|
|||
|
|
),
|
|||
|
|
required=True,
|
|||
|
|
max_length=1200,
|
|||
|
|
),
|
|||
|
|
],
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
await ctx.send_modal(modal)
|
|||
|
|
modal_ctx: ModalContext = await self.client.wait_for_modal(
|
|||
|
|
modal=modal, author=ctx.author
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
duration: datetime | None = (
|
|||
|
|
(datetime.now() + timedelta(minutes=duration)) if duration else None
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
title = modal_ctx.responses["title"]
|
|||
|
|
options: PollOptions = []
|
|||
|
|
for i, option in enumerate(
|
|||
|
|
modal_ctx.responses["options"].replace("-", "", 1).split("\n-")
|
|||
|
|
):
|
|||
|
|
if option == "":
|
|||
|
|
continue
|
|||
|
|
parts = option.split(":", 1)
|
|||
|
|
if len(parts) == 1:
|
|||
|
|
options.append((None, parts[0].strip()))
|
|||
|
|
else:
|
|||
|
|
options.append((parts[0].strip(), parts[1].strip()))
|
|||
|
|
|
|||
|
|
if len(options) > 10:
|
|||
|
|
await modal_ctx.send("You can only have up to 10 options.", ephemeral=True)
|
|||
|
|
return
|
|||
|
|
if len(options) < 2:
|
|||
|
|
await modal_ctx.send("You must have at least 2 options.", ephemeral=True)
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
poll_entry: PollsModel = PollsModel.create(
|
|||
|
|
guild_id=ctx.guild.id,
|
|||
|
|
author_id=ctx.author.id,
|
|||
|
|
title=title,
|
|||
|
|
options=json.dumps(options),
|
|||
|
|
no_options=len(options),
|
|||
|
|
multiple_choice=multiple_choice,
|
|||
|
|
expires=duration,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
buttons: List[Button] = []
|
|||
|
|
for i, option in enumerate(options):
|
|||
|
|
buttons.append(
|
|||
|
|
Button(
|
|||
|
|
emoji=(option[0] or num_to_emoji(i + 1)),
|
|||
|
|
style=ButtonStyles.PRIMARY,
|
|||
|
|
custom_id=f"poll-vote:{poll_entry.id}:{i}",
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
buttons.append(
|
|||
|
|
Button(
|
|||
|
|
label="Delete",
|
|||
|
|
style=ButtonStyles.DANGER,
|
|||
|
|
custom_id=f"poll-delete:{poll_entry.id}",
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
embed = generate_poll_embed(
|
|||
|
|
title,
|
|||
|
|
options,
|
|||
|
|
len(options) * [0],
|
|||
|
|
multiple_choice=multiple_choice,
|
|||
|
|
expires=duration,
|
|||
|
|
)
|
|||
|
|
poll_message: Message = await modal_ctx.send(
|
|||
|
|
embed=embed,
|
|||
|
|
components=spread_to_rows(*buttons),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
poll_entry.message_id = poll_message.id
|
|||
|
|
poll_entry.channel_id = poll_message.channel.id
|
|||
|
|
poll_entry.save()
|
|||
|
|
|
|||
|
|
@listen(events.Button)
|
|||
|
|
async def on_button(self, button: events.Button):
|
|||
|
|
ctx = button.context
|
|||
|
|
await ctx.defer(ephemeral=True)
|
|||
|
|
|
|||
|
|
if ctx.custom_id.startswith("poll-vote:"):
|
|||
|
|
poll_id, option_num = ctx.custom_id.split(":", 1)[1].split(":", 1)
|
|||
|
|
|
|||
|
|
poll_entry: PollsModel | None = PollsModel.get_or_none(
|
|||
|
|
guild_id=ctx.guild.id, id=poll_id
|
|||
|
|
)
|
|||
|
|
if not poll_entry:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
if poll_entry.expires and datetime.now() > poll_entry.expires:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
if not poll_entry.multiple_choice:
|
|||
|
|
votes_q: List[PollVotesModel] = PollVotesModel.select().where(
|
|||
|
|
PollVotesModel.poll_id == poll_id,
|
|||
|
|
PollVotesModel.user_id == ctx.author.id,
|
|||
|
|
)
|
|||
|
|
if votes_q.count() > 1:
|
|||
|
|
for vote in votes_q:
|
|||
|
|
vote.delete().execute()
|
|||
|
|
|
|||
|
|
PollVotesModel.create(
|
|||
|
|
poll_id=poll_id,
|
|||
|
|
user_id=ctx.author.id,
|
|||
|
|
option=option_num,
|
|||
|
|
)
|
|||
|
|
elif votes_q.count() == 1:
|
|||
|
|
if int(votes_q[0].option) == int(option_num):
|
|||
|
|
votes_q[0].delete_instance()
|
|||
|
|
await ctx.send("You have removed your vote.")
|
|||
|
|
else:
|
|||
|
|
votes_q[0].option = option_num
|
|||
|
|
votes_q[0].save()
|
|||
|
|
await ctx.send("You have changed your vote.")
|
|||
|
|
else:
|
|||
|
|
PollVotesModel.create(
|
|||
|
|
poll_id=poll_id,
|
|||
|
|
user_id=ctx.author.id,
|
|||
|
|
option=option_num,
|
|||
|
|
)
|
|||
|
|
await ctx.send("You have voted.")
|
|||
|
|
else:
|
|||
|
|
votes_q: List[PollVotesModel] = PollVotesModel.select().where(
|
|||
|
|
PollVotesModel.poll_id == poll_id,
|
|||
|
|
PollVotesModel.user_id == ctx.author.id,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
exists = False
|
|||
|
|
for vote in votes_q:
|
|||
|
|
if int(vote.option) == (option_num):
|
|||
|
|
exists = True
|
|||
|
|
vote.delete_instance()
|
|||
|
|
await ctx.send("You have removed your vote.")
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
if not exists:
|
|||
|
|
PollVotesModel.create(
|
|||
|
|
poll_id=poll_id,
|
|||
|
|
user_id=ctx.author.id,
|
|||
|
|
option=option_num,
|
|||
|
|
)
|
|||
|
|
await ctx.send("You have voted.")
|
|||
|
|
|
|||
|
|
votes_q: List[PollVotesModel] = (
|
|||
|
|
PollVotesModel.select(
|
|||
|
|
PollVotesModel.poll_id,
|
|||
|
|
PollVotesModel.option,
|
|||
|
|
fn.COUNT(PollVotesModel.option).alias("count"),
|
|||
|
|
)
|
|||
|
|
.where(PollVotesModel.poll_id == poll_id)
|
|||
|
|
.group_by(PollVotesModel.option)
|
|||
|
|
.order_by(PollVotesModel.option)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# This is such absolutely an awful way to do this. I'm sorry.
|
|||
|
|
# It's the cost of not adding the options in a separate table, I guess.
|
|||
|
|
# Anyway this just gets the votes for each option, and adds them to
|
|||
|
|
# the list `votes`. They're in the same order as the options.
|
|||
|
|
options: PollOptions = json.loads(poll_entry.options)
|
|||
|
|
votes = len(options) * [0]
|
|||
|
|
|
|||
|
|
for i, vote in enumerate(votes_q):
|
|||
|
|
votes[int(vote.option)] = vote.count
|
|||
|
|
|
|||
|
|
embed = generate_poll_embed(
|
|||
|
|
poll_entry.title,
|
|||
|
|
options,
|
|||
|
|
votes,
|
|||
|
|
multiple_choice=poll_entry.multiple_choice,
|
|||
|
|
expires=poll_entry.expires,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
await ctx.message.edit(embed=embed)
|
|||
|
|
|
|||
|
|
elif ctx.custom_id.startswith("poll-delete:"):
|
|||
|
|
poll_id = ctx.custom_id.split(":", 1)[1]
|
|||
|
|
|
|||
|
|
poll_entry: PollsModel | None = PollsModel.get_or_none(
|
|||
|
|
guild_id=ctx.guild.id, id=poll_id
|
|||
|
|
)
|
|||
|
|
if not poll_entry:
|
|||
|
|
await ctx.send("That poll doesn't exist.")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
if (
|
|||
|
|
not ctx.author.id == int(poll_entry.author_id)
|
|||
|
|
or not ctx.author.has_permission(Permissions.MANAGE_MESSAGES)
|
|||
|
|
):
|
|||
|
|
await ctx.send("You don't have permission to delete that poll.")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
poll_entry.delete_instance()
|
|||
|
|
PollVotesModel.delete().where(PollVotesModel.poll_id == poll_id).execute()
|
|||
|
|
await ctx.message.delete()
|
|||
|
|
await ctx.send("Poll deleted.")
|
|||
|
|
|
|||
|
|
@Task.create(IntervalTrigger(minutes=1))
|
|||
|
|
async def poll_expiry_check(self):
|
|||
|
|
logging.info("Checking for expired polls.")
|
|||
|
|
now = datetime.now()
|
|||
|
|
|
|||
|
|
polls_q: List[PollsModel] = PollsModel.select().where(PollsModel.expires < now)
|
|||
|
|
for poll_entry in polls_q:
|
|||
|
|
channel: GuildText = await self.client.fetch_channel(int(poll_entry.channel_id))
|
|||
|
|
if not channel:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
message: Message = await channel.fetch_message(int(poll_entry.message_id))
|
|||
|
|
if not message:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
await message.edit(components=[])
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
def setup(client: Client):
|
|||
|
|
PollsModel.create_table()
|
|||
|
|
PollVotesModel.create_table()
|
|||
|
|
Polls(client)
|
|||
|
|
logging.info("Polls extension loaded")
|