Skip to content

Commit

Permalink
Merge pull request #541 from allthingslinux/tess-remindme-rewrite
Browse files Browse the repository at this point in the history
  • Loading branch information
@kzndotsh
kzndotsh authored Sep 21, 2024
2 parents c3a6bd3 + 6e201ba commit 7b07a2e
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 114 deletions.
1 change: 1 addition & 0 deletions prisma/schema.prisma
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,7 @@ model Reminder {
reminder_expires_at DateTime
reminder_channel_id BigInt
reminder_user_id BigInt
reminder_sent Boolean @default(false)
guild_id BigInt
guild Guild @relation(fields: [guild_id], references: [guild_id])
Expand Down
143 changes: 30 additions & 113 deletions tux/cogs/utility/remindme.py
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
import asyncio
import contextlib
import datetime

import discord
from discord import app_commands
from discord.ext import commands
from discord.ext import commands, tasks
from loguru import logger

from prisma.models import Reminder
Expand All @@ -14,40 +13,26 @@
from tux.utils.functions import convert_to_seconds


def get_closest_reminder(reminders: list[Reminder]) -> Reminder | None:
"""
Check if there are any reminders and return the closest one.
Parameters
----------
reminders : list[Reminder]
A list of reminders to check.
Returns
-------
Reminder | None
The closest reminder or None if there are no reminders.
"""
return min(reminders, key=lambda x: x.reminder_expires_at) if reminders else None


class RemindMe(commands.Cog):
def __init__(self, bot: Tux) -> None:
self.bot = bot
self.db = DatabaseController().reminder
self.bot.loop.create_task(self.update())
self.check_reminders.start()

@tasks.loop(seconds=120)
async def check_reminders(self):
reminders = await self.db.get_unsent_reminders()

async def send_reminders(self, reminder: Reminder) -> None:
"""
Send the reminder to the user.
try:
for reminder in reminders:
await self.send_reminder(reminder)
await self.db.update_reminder_status(reminder.reminder_id, sent=True)
logger.debug(f'Status of reminder {reminder.reminder_id} updated to "sent".')

Parameters
----------
reminder : Reminder
The reminder object.
"""
except Exception as e:
logger.error(f"Error sending reminders: {e}")

async def send_reminder(self, reminder: Reminder) -> None:
user = self.bot.get_user(reminder.reminder_user_id)

if user is not None:
Expand All @@ -64,103 +49,37 @@ async def send_reminders(self, reminder: Reminder) -> None:
await user.send(embed=embed)

except discord.Forbidden:
# Send a message in the channel if the user has DMs closed
channel: discord.abc.GuildChannel | discord.Thread | discord.abc.PrivateChannel | None = (
self.bot.get_channel(reminder.reminder_channel_id)
)

if channel is not None and isinstance(
channel,
discord.TextChannel | discord.Thread | discord.VoiceChannel,
):
channel = self.bot.get_channel(reminder.reminder_channel_id)

if isinstance(channel, discord.TextChannel | discord.Thread | discord.VoiceChannel):
with contextlib.suppress(discord.Forbidden):
await channel.send(
content=f"{user.mention} Failed to DM you, sending in channel",
embed=embed,
)
return

else:
logger.error(
f"Failed to send reminder to {user.id}, DMs closed and channel not found.",
f"Failed to send reminder {reminder.reminder_id}, DMs closed and channel not found.",
)

else:
logger.error(f"Failed to send reminder to {reminder.reminder_user_id}, user not found.")

# Delete the reminder after sending
await self.db.delete_reminder_by_id(reminder.reminder_id)

# wait for a second so that the reminder is deleted before checking for more reminders
# who knows if this works, it seems to
await asyncio.sleep(1)

# Run update again to check if there are any more reminders
await self.update()

async def end_timer(self, reminder: Reminder) -> None:
"""
End the timer for the reminder.
Parameters
----------
reminder : Reminder
The reminder object.
"""

# Wait until the reminder expires
await discord.utils.sleep_until(reminder.reminder_expires_at)
await self.send_reminders(reminder)

async def update(self) -> None:
"""
Update the reminders
Check if there are any reminders and send the closest one.
"""

try:
# Get all reminders
reminders = await self.db.get_all_reminders()
# Get the closest reminder
closest_reminder = get_closest_reminder(reminders)

except Exception as e:
logger.error(f"Error getting reminders: {e}")
return

# If there are no reminders, return
if closest_reminder is None:
return

# Check if it's expired
if closest_reminder.reminder_expires_at < datetime.datetime.now(datetime.UTC):
await self.send_reminders(closest_reminder)
return
logger.error(
f"Failed to send reminder {reminder.reminder_id}, user with ID {reminder.reminder_user_id} not found.",
)

# Create a task to wait until the reminder expires
self.bot.loop.create_task(self.end_timer(closest_reminder))
@check_reminders.before_loop
async def before_check_reminders(self):
await self.bot.wait_until_ready()

@app_commands.command(
name="remindme",
description="Reminds you after a certain amount of time.",
)
async def remindme(self, interaction: discord.Interaction, time: str, *, reminder: str) -> None:
"""
Set a reminder for a certain amount of time.
Parameters
----------
interaction : discord.Interaction
The discord interaction object.
time : str
Time in the format `[number][M/w/d/h/m/s]`.
reminder : str
Reminder content.
"""

seconds = convert_to_seconds(time)

# Check if the time is valid (this is set to 0 if the time is invalid via convert_to_seconds)
if seconds == 0:
await interaction.response.send_message(
"Invalid time format. Please use the format `[number][M/w/d/h/m/s]`.",
Expand All @@ -169,13 +88,13 @@ async def remindme(self, interaction: discord.Interaction, time: str, *, reminde
)
return

seconds = datetime.datetime.now(datetime.UTC) + datetime.timedelta(seconds=seconds)
expires_at = datetime.datetime.now(datetime.UTC) + datetime.timedelta(seconds=seconds)

try:
await self.db.insert_reminder(
reminder_user_id=interaction.user.id,
reminder_content=reminder,
reminder_expires_at=seconds,
reminder_expires_at=expires_at,
reminder_channel_id=interaction.channel_id or 0,
guild_id=interaction.guild_id or 0,
)
Expand All @@ -186,12 +105,13 @@ async def remindme(self, interaction: discord.Interaction, time: str, *, reminde
user_name=interaction.user.name,
user_display_avatar=interaction.user.display_avatar.url,
title="Reminder Set",
description=f"Reminder set for <t:{int(seconds.timestamp())}:f>.",
description=f"Reminder set for <t:{int(expires_at.timestamp())}:f>.",
)

embed.add_field(
name="Note",
value="If you have DMs closed, the reminder may not reach you. We will attempt to send it in this channel instead, however it is not guaranteed.",
value="- If you have DMs closed, we will attempt to send it in this channel instead.\n"
"- The reminder may be delayed by up to 120 seconds due to the way Tux works.",
)

except Exception as e:
Expand All @@ -207,9 +127,6 @@ async def remindme(self, interaction: discord.Interaction, time: str, *, reminde

await interaction.response.send_message(embed=embed, ephemeral=True)

# Run update again to check if this reminder is the closest
await self.update()


async def setup(bot: Tux) -> None:
await bot.add_cog(RemindMe(bot))
23 changes: 22 additions & 1 deletion tux/database/controllers/reminder.py
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from datetime import datetime
from datetime import UTC, datetime

from prisma.models import Guild, Reminder
from tux.database.client import db
Expand All @@ -21,6 +21,10 @@ async def get_all_reminders(self) -> list[Reminder]:
async def get_reminder_by_id(self, reminder_id: int) -> Reminder | None:
return await self.table.find_first(where={"reminder_id": reminder_id})

async def get_unsent_reminders(self) -> list[Reminder]:
now = datetime.now(UTC)
return await self.table.find_many(where={"reminder_sent": False, "reminder_expires_at": {"lte": now}})

async def insert_reminder(
self,
reminder_user_id: int,
Expand All @@ -38,6 +42,7 @@ async def insert_reminder(
"reminder_expires_at": reminder_expires_at,
"reminder_channel_id": reminder_channel_id,
"guild_id": guild_id,
"reminder_sent": False,
},
)

Expand All @@ -53,3 +58,19 @@ async def update_reminder_by_id(
where={"reminder_id": reminder_id},
data={"reminder_content": reminder_content},
)

async def update_reminder_status(self, reminder_id: int, sent: bool = True) -> None:
"""
Update the status of a reminder. This sets the value "reminder_sent" to True by default.
Parameters
----------
reminder_id : int
The ID of the reminder to update.
sent : bool
The new status of the reminder.
"""
await self.table.update(
where={"reminder_id": reminder_id},
data={"reminder_sent": sent},
)

0 comments on commit 7b07a2e

Please sign in to comment.