2022-10-15 03:35:42 +02:00
|
|
|
|
# pylint: disable=not-an-iterable
|
|
|
|
|
# pylint: disable=unsubscriptable-object
|
|
|
|
|
# pylint: disable=logging-fstring-interpolation
|
2022-08-05 13:19:08 +02:00
|
|
|
|
from datetime import datetime, timedelta
|
|
|
|
|
import logging
|
|
|
|
|
import json
|
2022-10-15 03:35:42 +02:00
|
|
|
|
from typing import List, Tuple
|
2022-08-05 13:19:08 +02:00
|
|
|
|
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,
|
2022-08-06 02:01:29 +02:00
|
|
|
|
PartialEmoji,
|
2022-08-05 13:19:08 +02:00
|
|
|
|
)
|
|
|
|
|
from naff.api import events
|
2022-08-06 02:01:29 +02:00
|
|
|
|
from naff.client.errors import HTTPException
|
2022-08-05 13:19:08 +02:00
|
|
|
|
|
|
|
|
|
from peewee import fn
|
2022-10-15 03:35:42 +02:00
|
|
|
|
from database import Polls as PollsModel, PollVotes as PollVotesModel
|
2022-08-05 13:19:08 +02:00
|
|
|
|
|
|
|
|
|
PollOptions = List[Tuple[str | None, str]]
|
|
|
|
|
|
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
def datetime_to_discord_relative_time(dt: datetime) -> str:
|
|
|
|
|
"""Create a Discord relative time text from a datetime object."""
|
2022-08-05 13:19:08 +02:00
|
|
|
|
t = dt.strftime("%s")
|
|
|
|
|
return f"<t:{int(t)}:R>"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def generate_bar(num: int, total: int, length: int = 10) -> str:
|
2022-08-11 01:39:15 +02:00
|
|
|
|
"""Create a bar graph from a number and a total.
|
2022-10-15 03:35:42 +02:00
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
Parameters
|
|
|
|
|
----------
|
|
|
|
|
num : int
|
|
|
|
|
The current amount.
|
|
|
|
|
total : int
|
|
|
|
|
The total amount.
|
|
|
|
|
length : int
|
|
|
|
|
The amount of characters to use for the bar.
|
|
|
|
|
"""
|
2022-08-05 13:19:08 +02:00
|
|
|
|
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:
|
2022-08-11 01:39:15 +02:00
|
|
|
|
"""Create a poll embed from a title, options and votes.
|
2022-10-15 03:35:42 +02:00
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
Parameters
|
|
|
|
|
----------
|
|
|
|
|
title : str
|
|
|
|
|
The title of the poll.
|
|
|
|
|
options : PollOptions
|
|
|
|
|
The options of the poll.
|
|
|
|
|
votes : List[int]
|
|
|
|
|
The votes of the poll, in the same order as the options.
|
|
|
|
|
multiple_choice : bool, optional
|
|
|
|
|
Whether the poll is multiple choice. Defaults to False.
|
|
|
|
|
expires : datetime, optional
|
|
|
|
|
The time at which the poll expires. Defaults to None.
|
|
|
|
|
"""
|
2022-08-05 13:19:08 +02:00
|
|
|
|
data = []
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# \u2022 is a bullet point
|
2022-08-05 13:19:08 +02:00
|
|
|
|
if multiple_choice:
|
|
|
|
|
data.append("\u2022 Multiple choice")
|
|
|
|
|
if expires:
|
2022-08-11 01:39:15 +02:00
|
|
|
|
data.append(f"\u2022 Expiry: {datetime_to_discord_relative_time(expires)}")
|
2022-08-05 13:19:08 +02:00
|
|
|
|
|
|
|
|
|
embed = Embed(
|
|
|
|
|
title=title,
|
|
|
|
|
description=("\n".join(data) if data else None),
|
2022-08-06 02:01:29 +02:00
|
|
|
|
)
|
2022-08-05 13:19:08 +02:00
|
|
|
|
sum_votes = sum(votes)
|
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# Add a field for each option, with the vote count and a bar graph showing the percentage.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
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
|
|
|
|
|
|
|
|
|
|
|
2022-10-15 03:35:42 +02:00
|
|
|
|
def num_to_emoji(num: int) -> str: # pylint: disable=too-many-return-statements
|
2022-08-05 13:19:08 +02:00
|
|
|
|
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 _:
|
2022-10-15 03:35:42 +02:00
|
|
|
|
raise ValueError("Invalid number: `num` must be 0 <= num <= 10.")
|
2022-08-05 13:19:08 +02:00
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class Polls(Extension):
|
|
|
|
|
def __init__(self, client: Client):
|
|
|
|
|
self.client = client
|
|
|
|
|
|
|
|
|
|
@listen(events.Ready)
|
|
|
|
|
async def on_ready(self):
|
|
|
|
|
await self.poll_expiry_check()
|
2022-10-15 03:35:42 +02:00
|
|
|
|
self.poll_expiry_check.start() # pylint: disable=no-member
|
2022-08-05 13:19:08 +02:00
|
|
|
|
|
|
|
|
|
@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,
|
|
|
|
|
)
|
2022-10-15 03:35:42 +02:00
|
|
|
|
async def create_poll( # pylint: disable=too-many-locals
|
2022-08-05 13:19:08 +02:00
|
|
|
|
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,
|
|
|
|
|
),
|
|
|
|
|
],
|
|
|
|
|
)
|
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# Display the modal and wait for the user to submit.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
await ctx.send_modal(modal)
|
|
|
|
|
modal_ctx: ModalContext = await self.client.wait_for_modal(
|
|
|
|
|
modal=modal, author=ctx.author
|
|
|
|
|
)
|
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# If the user set a duration for the poll, create a datetime for the time in the
|
|
|
|
|
# future when the poll expires.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
duration: datetime | None = (
|
|
|
|
|
(datetime.now() + timedelta(minutes=duration)) if duration else None
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
title = modal_ctx.responses["title"]
|
|
|
|
|
options: PollOptions = []
|
2022-08-11 01:39:15 +02:00
|
|
|
|
|
|
|
|
|
# Get each option from the poll options.
|
|
|
|
|
# We're stripping the first occurance of a dash, as it otherwise would be included.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
for i, option in enumerate(
|
|
|
|
|
modal_ctx.responses["options"].replace("-", "", 1).split("\n-")
|
|
|
|
|
):
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# If the option is empty, skip it.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
if option == "":
|
|
|
|
|
continue
|
2022-08-11 01:39:15 +02:00
|
|
|
|
|
|
|
|
|
# Check if the option contains an emoji.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
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,
|
|
|
|
|
)
|
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# Create vote buttons for each option.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
buttons: List[Button] = []
|
|
|
|
|
for i, option in enumerate(options):
|
2022-08-06 02:01:29 +02:00
|
|
|
|
try:
|
|
|
|
|
emoji = PartialEmoji.from_str((option[0] or num_to_emoji(i + 1)))
|
|
|
|
|
except ValueError:
|
|
|
|
|
emoji = num_to_emoji(i + 1)
|
2022-08-05 13:19:08 +02:00
|
|
|
|
buttons.append(
|
|
|
|
|
Button(
|
2022-08-06 02:01:29 +02:00
|
|
|
|
emoji=emoji,
|
2022-08-05 13:19:08 +02:00
|
|
|
|
style=ButtonStyles.PRIMARY,
|
|
|
|
|
custom_id=f"poll-vote:{poll_entry.id}:{i}",
|
|
|
|
|
)
|
|
|
|
|
)
|
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# Create a button to allow locking the poll.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
buttons.append(
|
|
|
|
|
Button(
|
2022-08-06 02:01:29 +02:00
|
|
|
|
emoji="🔒",
|
|
|
|
|
label="Lock",
|
2022-08-05 13:19:08 +02:00
|
|
|
|
style=ButtonStyles.DANGER,
|
2022-08-06 02:01:29 +02:00
|
|
|
|
custom_id=f"poll-lock:{poll_entry.id}",
|
2022-08-05 13:19:08 +02:00
|
|
|
|
)
|
|
|
|
|
)
|
|
|
|
|
embed = generate_poll_embed(
|
|
|
|
|
title,
|
|
|
|
|
options,
|
|
|
|
|
len(options) * [0],
|
|
|
|
|
multiple_choice=multiple_choice,
|
|
|
|
|
expires=duration,
|
|
|
|
|
)
|
2022-08-06 02:01:29 +02:00
|
|
|
|
try:
|
|
|
|
|
poll_message: Message = await modal_ctx.send(
|
|
|
|
|
embed=embed,
|
|
|
|
|
components=spread_to_rows(*buttons),
|
|
|
|
|
)
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# Naive error handling.
|
2022-08-06 02:01:29 +02:00
|
|
|
|
except HTTPException as e:
|
|
|
|
|
logging.error(f"Error sending poll: {e}")
|
|
|
|
|
await modal_ctx.send(
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# TODO: This should probably also include a sample for poll options.
|
2022-08-06 02:01:29 +02:00
|
|
|
|
"Error during poll creation. NB: You can not use server-specific emojis",
|
|
|
|
|
ephemeral=True,
|
|
|
|
|
)
|
|
|
|
|
return
|
2022-08-05 13:19:08 +02:00
|
|
|
|
|
|
|
|
|
poll_entry.message_id = poll_message.id
|
|
|
|
|
poll_entry.channel_id = poll_message.channel.id
|
|
|
|
|
poll_entry.save()
|
|
|
|
|
|
|
|
|
|
@listen(events.Button)
|
2022-10-15 03:35:42 +02:00
|
|
|
|
async def on_button(self, button: events.Button): #pylint: disable=too-many-branches,too-many-statements
|
2022-08-05 13:19:08 +02:00
|
|
|
|
ctx = button.context
|
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# Ensure that the pressed button is a vote button.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
if ctx.custom_id.startswith("poll-vote:"):
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# Get the poll ID and the option index.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
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
|
|
|
|
|
|
2022-10-15 02:51:21 +02:00
|
|
|
|
await ctx.defer(ephemeral=True)
|
2022-08-05 13:19:08 +02:00
|
|
|
|
if not poll_entry.multiple_choice:
|
|
|
|
|
votes_q: List[PollVotesModel] = PollVotesModel.select().where(
|
|
|
|
|
PollVotesModel.poll_id == poll_id,
|
|
|
|
|
PollVotesModel.user_id == ctx.author.id,
|
|
|
|
|
)
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# If the user somehow already has more than one vote, delete them.
|
|
|
|
|
# This should never happen, but just in case.
|
|
|
|
|
# Then, add the new vote.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
if votes_q.count() > 1:
|
|
|
|
|
for vote in votes_q:
|
2022-08-11 01:39:15 +02:00
|
|
|
|
vote.delete_instance()
|
2022-08-05 13:19:08 +02:00
|
|
|
|
|
|
|
|
|
PollVotesModel.create(
|
|
|
|
|
poll_id=poll_id,
|
|
|
|
|
user_id=ctx.author.id,
|
|
|
|
|
option=option_num,
|
|
|
|
|
)
|
|
|
|
|
elif votes_q.count() == 1:
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# If the vote is the current vote, delete it.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
if int(votes_q[0].option) == int(option_num):
|
|
|
|
|
votes_q[0].delete_instance()
|
|
|
|
|
await ctx.send("You have removed your vote.")
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# If it's not the current vote, change the vote to the new one.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
else:
|
|
|
|
|
votes_q[0].option = option_num
|
|
|
|
|
votes_q[0].save()
|
|
|
|
|
await ctx.send("You have changed your vote.")
|
2022-08-11 01:39:15 +02:00
|
|
|
|
#If the user has no votes, add a new vote.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
else:
|
|
|
|
|
PollVotesModel.create(
|
|
|
|
|
poll_id=poll_id,
|
|
|
|
|
user_id=ctx.author.id,
|
|
|
|
|
option=option_num,
|
|
|
|
|
)
|
|
|
|
|
await ctx.send("You have voted.")
|
2022-08-11 01:39:15 +02:00
|
|
|
|
|
|
|
|
|
# If the poll is multiple choice
|
2022-08-05 13:19:08 +02:00
|
|
|
|
else:
|
|
|
|
|
votes_q: List[PollVotesModel] = PollVotesModel.select().where(
|
|
|
|
|
PollVotesModel.poll_id == poll_id,
|
|
|
|
|
PollVotesModel.user_id == ctx.author.id,
|
|
|
|
|
)
|
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# If the user has already voted for this option, remove their vote.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
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
|
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# If the user has not voted for this option, add a new vote.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
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]
|
|
|
|
|
|
2022-10-15 03:35:42 +02:00
|
|
|
|
for vote in votes_q:
|
2022-08-05 13:19:08 +02:00
|
|
|
|
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,
|
|
|
|
|
)
|
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# Edit the message with the updated information.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
await ctx.message.edit(embed=embed)
|
2022-08-06 02:01:29 +02:00
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# If the "lock poll" button is pressed, lock the poll.
|
2022-08-06 02:01:29 +02:00
|
|
|
|
elif ctx.custom_id.startswith("poll-lock:"):
|
2022-08-05 13:19:08 +02:00
|
|
|
|
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
|
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# Ensure that the user is the poll creator, or can manage messages.
|
2022-08-06 02:01:29 +02:00
|
|
|
|
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 lock that poll.")
|
2022-08-05 13:19:08 +02:00
|
|
|
|
return
|
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# Set the poll to be expired to lock it.
|
2022-08-06 02:01:29 +02:00
|
|
|
|
poll_entry.expires = datetime.now() - timedelta(minutes=1)
|
|
|
|
|
poll_entry.save()
|
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# Force the "poll expiry check" task to run.
|
2022-08-06 02:01:29 +02:00
|
|
|
|
await self.poll_expiry_check()
|
|
|
|
|
await ctx.send("Poll locked.")
|
2022-08-05 13:19:08 +02:00
|
|
|
|
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# A task that runs each minute to check for expired polls.
|
2022-08-05 13:19:08 +02:00
|
|
|
|
@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:
|
2022-08-06 02:01:29 +02:00
|
|
|
|
channel: GuildText = await self.client.fetch_channel(
|
|
|
|
|
int(poll_entry.channel_id)
|
|
|
|
|
)
|
2022-08-05 13:19:08 +02:00
|
|
|
|
if not channel:
|
|
|
|
|
continue
|
|
|
|
|
|
|
|
|
|
message: Message = await channel.fetch_message(int(poll_entry.message_id))
|
|
|
|
|
if not message:
|
|
|
|
|
continue
|
|
|
|
|
|
|
|
|
|
await message.edit(components=[])
|
2022-08-11 01:39:15 +02:00
|
|
|
|
# Delete associated database entries, as they will no longer be updated.
|
|
|
|
|
PollVotesModel.delete().where(PollVotesModel.poll_id == poll_entry.id).execute()
|
|
|
|
|
poll_entry.delete_instance()
|
2022-10-15 03:35:42 +02:00
|
|
|
|
|
2022-08-05 13:19:08 +02:00
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def setup(client: Client):
|
|
|
|
|
PollsModel.create_table()
|
|
|
|
|
PollVotesModel.create_table()
|
|
|
|
|
Polls(client)
|
|
|
|
|
logging.info("Polls extension loaded")
|