Files
croccybot/src/modules/reminders/cog.py

631 lines
22 KiB
Python

"""
Max 25 reminders (propagating Discord restriction)
/reminders show
-- Widget which displays and allows removing reminders.
-- Points to /remindme for setting
/reminders clear
/reminders remove <reminder: acmpl>
-- Can we autocomplete an integer field?
/remindme at <time: time> <repeat every: acmpl str> <reminder: str>
/remindme in <days: int> <hours: int> <minutes: int> <repeat every: acmpl str> <reminder: str>
"""
from typing import Optional
import datetime as dt
from cachetools import TTLCache
import discord
from discord.ext import commands as cmds
from discord import app_commands as appcmds
from discord.app_commands import Transform
from dateutil.parser import parse, ParserError
from data.queries import ORDER
from meta import LionBot, LionCog, LionContext
from meta.errors import UserInputError
from meta.app import shard_talk, appname_from_shard
from meta.logger import log_wrap, set_logging_context
from babel import ctx_translator, ctx_locale
from utils.lib import parse_duration, utc_now, strfdur, error_embed, check_dm
from utils.monitor import TaskMonitor
from utils.transformers import DurationTransformer
from utils.ui import AButton, AsComponents
from utils.ratelimits import Bucket
from . import babel, logger
from .data import ReminderData
from .ui import ReminderList
_, _p, _np = babel._, babel._p, babel._np
class ReminderMonitor(TaskMonitor[int]):
...
class Reminders(LionCog):
def __init__(self, bot: LionBot):
self.bot = bot
self.data = bot.db.load_registry(ReminderData())
# Whether this process should handle reminder execution
self.executor = (self.bot.shard_id == 0)
self.executor_name = appname_from_shard(0)
if self.executor:
self.monitor: Optional[ReminderMonitor] = ReminderMonitor(
executor=self.execute_reminder,
bucket=Bucket(5, 10)
)
else:
self.monitor = None
self.talk_reload = shard_talk.register_route('reload_reminders')(self.reload_reminders)
self.talk_schedule = shard_talk.register_route('schedule_reminders')(self.schedule_reminders)
self.talk_cancel = shard_talk.register_route('cancel_reminders')(self.cancel_reminders)
# Short term userid -> list[Reminder] cache, mainly for autocomplete
self._user_reminder_cache: TTLCache[int, list[ReminderData.Reminder]] = TTLCache(1000, ttl=60)
self._active_reminderlists: dict[int, ReminderList] = {}
async def cog_load(self):
await self.data.init()
if self.executor and self.bot.is_ready():
await self.on_ready()
@LionCog.listener()
async def on_ready(self):
if self.executor:
if self.monitor and self.monitor._monitor_task:
self.monitor._monitor_task.cancel()
# Attach and populate the reminder monitor
self.monitor = ReminderMonitor(executor=self.execute_reminder, bucket=Bucket(5, 10))
await self.reload_reminders()
# Start firing reminders
self.monitor.start()
# ----- Cog API -----
async def create_reminder(
self,
userid: int, remind_at: dt.datetime, content: str,
message_link: Optional[str] = None,
interval: Optional[int] = None,
created_at: Optional[dt.datetime] = None,
) -> ReminderData.Reminder:
"""
Create and schedule a new reminder from user-entered data.
Raises UserInputError if the requested parameters are invalid.
"""
now = utc_now()
if remind_at <= now:
t = self.bot.translator.t
raise UserInputError(
t(_p(
'create_reminder|error:past',
"The provided reminder time {timestamp} is in the past!"
)).format(timestamp=discord.utils.format_dt(remind_at))
)
if interval is not None and interval < 600:
t = self.bot.translator.t
raise UserInputError(
t(_p(
'create_reminder|error:too_fast',
"You cannot set a repeating reminder with a period less than 10 minutes."
))
)
existing = await self.data.Reminder.fetch_where(userid=userid)
if len(existing) >= 25:
t = self.bot.translator.t
raise UserInputError(
t(_p(
'create_reminder|error:too_many',
"Sorry, you have reached the maximum of `25` reminders."
))
)
user = self.bot.get_user(userid)
if not user:
user = await self.bot.fetch_user(userid)
if not user:
raise ValueError(f"Target user {userid} does not exist.")
can_dm = await check_dm(user)
if not can_dm:
t = self.bot.translator.t
raise UserInputError(
t(_p(
'create_reminder|error:cannot_dm',
"I cannot direct message you! Do you have me blocked or direct messages closed?"
))
)
created_at = created_at or now
# Passes validation, actually create
reminder = await self.data.Reminder.create(
userid=userid,
remind_at=remind_at,
content=content,
message_link=message_link,
interval=interval,
created_at=created_at,
)
# Schedule from executor
await self.talk_schedule(reminder.reminderid).send(self.executor_name, wait_for_reply=False)
# Dispatch reminder update
await self.dispatch_update_for(userid)
# Return fresh reminder
return reminder
async def parse_time_static(self, timestr, timezone):
timestr = timestr.strip()
default = dt.datetime.now(tz=timezone).replace(hour=0, minute=0, second=0, microsecond=0)
if not timestr:
return default
try:
ts = parse(timestr, fuzzy=True, default=default)
except ParserError:
t = self.bot.translator.t
raise UserInputError(
t(_p(
'parse_timestamp|error:parse',
"Could not parse `{given}` as a valid reminder time. "
"Try entering the time in the form `HH:MM` or `YYYY-MM-DD HH:MM`."
)).format(given=timestr)
)
return ts
async def get_reminders_for(self, userid: int):
"""
Retrieve a list of reminders for the given userid, using the cache.
"""
reminders = self._user_reminder_cache.get(userid, None)
if reminders is None:
reminders = await self.data.Reminder.fetch_where(
userid=userid
).order_by(self.data.Reminder.created_at, ORDER.ASC)
self._user_reminder_cache[userid] = reminders
return reminders
async def dispatch_update_for(self, userid: int):
"""
Announce that the given user's reminders have changed.
This triggers update of the cog reminder cache, and a reload of any active reminder list UIs.
"""
self._user_reminder_cache.pop(userid, None)
if userid in self._active_reminderlists:
await self._active_reminderlists[userid].refresh()
async def reload_reminders(self):
"""
Refresh reminder data and reminder tasks.
"""
if not self.executor:
raise ValueError("Only the executor shard can reload reminders!")
# Load all reminder tasks
reminders = await self.data.Reminder.fetch_where(
self.data.Reminder.remind_at > utc_now(),
failed=None
)
tasks = [(r.reminderid, r.timestamp) for r in reminders]
self.monitor.set_tasks(*tasks)
logger.info(
f"Reloaded ReminderMonitor with {len(tasks)} active reminders."
)
async def cancel_reminders(self, *reminderids):
"""
ShardTalk Route.
Cancel the given reminderids in the ReminderMonitor.
"""
if not self.executor:
raise ValueError("Only the executor shard can cancel scheduled reminders!")
# If we are the executor shard, we know the monitor is loaded
# If reminders have not yet been loaded, cancelling is a no-op
# Since reminder loading is synchronous, we cannot get in a race state with loading
self.monitor.cancel_tasks(*reminderids)
logger.debug(
f"Cancelled reminders: {reminderids}",
)
async def schedule_reminders(self, *reminderids):
"""
ShardTalk Route.
Schedule the given new reminderids in the ReminderMonitor.
"""
if not self.executor:
raise ValueError("Only the executor shard can schedule reminders!")
# We refetch here to make sure the reminders actually exist
reminders = await self.data.Reminder.fetch_where(reminderid=reminderids)
self.monitor.schedule_tasks(*((reminder.reminderid, reminder.timestamp) for reminder in reminders))
logger.debug(
f"Scheduled new reminders: {tuple(reminder.reminderid for reminder in reminders)}",
)
@log_wrap(action="Send Reminder")
async def execute_reminder(self, reminderid):
"""
Send the reminder with the given reminderid.
This should in general only be executed from the executor shard,
through a ReminderMonitor instance.
"""
set_logging_context(context=f"rid: {reminderid}")
reminder = await self.data.Reminder.fetch(reminderid)
if reminder is None:
logger.warning(
f"Attempted to execute a reminder <rid: {reminderid}> that no longer exists!"
)
return
try:
# Try and find the user
userid = reminder.userid
if not (user := self.bot.get_user(userid)):
user = await self.bot.fetch_user(userid)
# Set the locale variables
locale = await self.bot.get_cog('BabelCog').get_user_locale(userid)
ctx_locale.set(locale)
ctx_translator.set(self.bot.translator)
# Build the embed
embed = reminder.embed
# Attempt to send to user
# TODO: Consider adding a View to this, for cancelling a repeated reminder or showing reminders
await user.send(embed=embed)
# Update the data as required
if reminder.interval:
now = utc_now()
# Use original reminder time to calculate repeat, avoiding drift
next_time = reminder.remind_at + dt.timedelta(seconds=reminder.interval)
# Skip any expired repeats, to avoid spamming requests after downtime
# TODO: Is this actually dst safe?
while next_time.timestamp() <= now.timestamp():
next_time = next_time + dt.timedelta(seconds=reminder.interval)
await reminder.update(remind_at=next_time)
self.monitor.schedule_task(reminder.reminderid, reminder.timestamp)
logger.debug(
f"Executed reminder <rid: {reminder.reminderid}> and scheduled repeat at {next_time}."
)
else:
await reminder.delete()
logger.debug(
f"Executed reminder <rid: {reminder.reminderid}>."
)
except discord.HTTPException as e:
await reminder.update(failed=True)
logger.debug(
f"Reminder <rid: {reminder.reminderid}> could not be sent: {e.text}",
)
except Exception:
await reminder.update(failed=True)
logger.exception(
f"Reminder <rid: {reminder.reminderid}> failed for an unknown reason!"
)
finally:
# Dispatch for analytics
self.bot.dispatch('reminder_sent', reminder)
@cmds.hybrid_command(
name=_p('cmd:reminders', "reminders"),
description=_p(
'cmd:reminders|desc',
"View and set your reminders."
)
)
@appcmds.guild_only
async def cmd_reminders(self, ctx: LionContext):
"""
Display the reminder widget for this user.
"""
if not ctx.interaction:
return
if ctx.author.id in self._active_reminderlists:
await self._active_reminderlists[ctx.author.id].quit()
ui = ReminderList(self.bot, ctx.author)
try:
self._active_reminderlists[ctx.author.id] = ui
await ui.run(ctx.interaction, ephemeral=True)
await ui.wait()
finally:
self._active_reminderlists.pop(ctx.author.id, None)
@cmds.hybrid_group(
name=_p('cmd:remindme', "remindme"),
description=_p('cmd:remindme|desc', "View and set task reminders."),
)
@appcmds.guild_only
async def remindme_group(self, ctx: LionContext):
# Base command group for scheduling reminders.
pass
@remindme_group.command(
name=_p('cmd:reminders_cancel', "cancel"),
description=_p(
'cmd:reminders_cancel|desc',
"Cancel a single reminder. Use /reminders to clear or cancel multiple reminders."
)
)
@appcmds.rename(
reminder=_p('cmd:reminders_cancel|param:reminder', 'reminder')
)
@appcmds.describe(
reminder=_p(
'cmd:reminders_cancel|param:reminder|desc',
"Start typing, then select a reminder to cancel."
)
)
async def cmd_reminders_cancel(self, ctx: LionContext, reminder: str):
# No help string
"""
Cancel a previously scheduled reminder.
Autocomplete lets the user select their reminder by number or truncated content.
Need to handle the case where reminderid is that truncated content.
"""
t = self.bot.translator.t
reminders = await self.get_reminders_for(ctx.author.id)
# Guard against no reminders
if not reminders:
await ctx.error_reply(
t(_p(
'cmd:reminders_cancel|error:no_reminders',
"There are no reminders to cancel!"
))
)
return
# Now attempt to parse reminder input
if reminder.startswith('rid:') and reminder[4:].isdigit():
# Assume reminderid, probably selected through autocomplete
rid = int(reminder[4:])
rem = next((rem for rem in reminders if rem.reminderid == rid), None)
elif reminder.strip('[] ').isdigit():
# Assume user reminder index
# Not strictly threadsafe, but should be okay 90% of the time
lid = int(reminder)
rem = next((rem for i, rem in enumerate(reminders, start=1) if i == lid), None)
else:
# Assume partial string from a reminder
partial = reminder
rem = next((rem for rem in reminders if partial in rem.content), None)
if rem is None:
await ctx.error_reply(
t(_p(
'cmd:reminders_cancel|error:no_match',
"I am not sure which reminder you want to cancel. "
"Please try again, selecting a reminder from the list of choices."
))
)
return
# At this point we have a valid reminder to cancel
await rem.delete()
await self.talk_cancel(rem.reminderid).send(self.executor_name, wait_for_reply=False)
await ctx.reply(
embed=discord.Embed(
description=t(_p(
'cmd:reminders_cancel|embed:success|desc',
"Reminder successfully cancelled."
)),
colour=discord.Colour.brand_green()
),
ephemeral=True
)
await self.dispatch_update_for(ctx.author.id)
@cmd_reminders_cancel.autocomplete('reminder')
async def cmd_reminders_cancel_acmpl_reminderid(self, interaction: discord.Interaction, partial: str):
t = self.bot.translator.t
reminders = await self.get_reminders_for(interaction.user.id)
if not reminders:
# Nothing to cancel case
name = t(_p(
'cmd:reminders_cancel|acmpl:reminder|error:no_reminders',
"There are no reminders to cancel!"
))
value = 'None'
choices = [
appcmds.Choice(name=name[:100], value=value)
]
else:
# Build list of reminder strings
strings = []
for pos, reminder in enumerate(reminders, start=1):
strings.append(
(f"[{pos}] {reminder.content}", reminder)
)
# Extract matches
matches = [string for string in strings if partial.lower() in string[0].lower()]
if matches:
# Build list of valid choices
choices = [
appcmds.Choice(
name=string[0][:100],
value=f"rid:{string[1].reminderid}"
)
for string in matches
]
else:
choices = [
appcmds.Choice(
name=t(_p(
'cmd:reminders_cancel|acmpl:reminder|error:no_matches',
"You do not have any reminders matching \"{partial}\""
)).format(partial=partial)[:100],
value=partial
)
]
return choices
@remindme_group.command(
name=_p('cmd:remindme_at', "at"),
description=_p(
'cmd:remindme_at|desc',
"Schedule a reminder for a particular time."
)
)
@appcmds.rename(
time=_p('cmd:remindme_at|param:time', "time"),
reminder=_p('cmd:remindme_at|param:reminder', "reminder"),
every=_p('cmd:remindme_at|param:every', "repeat_every"),
)
@appcmds.describe(
time=_p(
'cmd:remindme_at|param:time|desc',
"When you want to be reminded. (E.g. `4pm` or `16:00`)."
),
reminder=_p(
'cmd:remindme_at|param:reminder|desc',
"What should the reminder be?"
),
every=_p(
'cmd:remindme_at|param:every|desc',
"How often to repeat this reminder."
)
)
async def cmd_remindme_at(
self,
ctx: LionContext,
time: appcmds.Range[str, 1, 100],
reminder: appcmds.Range[str, 1, 2000],
every: Optional[Transform[int, DurationTransformer(60)]] = None
):
t = self.bot.translator.t
try:
timezone = ctx.lmember.timezone
remind_at = await self.parse_time_static(time, timezone)
reminder = await self.create_reminder(
userid=ctx.author.id,
remind_at=remind_at,
content=reminder,
message_link=ctx.message.jump_url,
interval=every,
)
embed = reminder.set_response
except UserInputError as e:
embed = discord.Embed(
title=t(_p(
'cmd:remindme_at|error|title',
"Could not create reminder!"
)),
description=e.msg,
colour=discord.Colour.brand_red()
)
await ctx.reply(
embed=embed,
ephemeral=True
)
@cmd_remindme_at.autocomplete('time')
async def cmd_remindme_at_acmpl_time(self, interaction: discord.Interaction, partial: str):
if interaction.guild:
lmember = await self.bot.core.lions.fetch_member(interaction.guild.id, interaction.user.id)
timezone = lmember.timezone
else:
luser = await self.bot.core.lions.fetch_user(interaction.user.id)
timezone = luser.timezone
t = self.bot.translator.t
try:
timestamp = await self.parse_time_static(partial, timezone)
choice = appcmds.Choice(
name=timestamp.strftime('%Y-%m-%d %H:%M'),
value=partial
)
except UserInputError:
choice = appcmds.Choice(
name=t(_p(
'cmd:remindme_at|acmpl:time|error:parse',
"Cannot parse \"{partial}\" as a time. Try the format HH:MM or YYYY-MM-DD HH:MM"
)).format(partial=partial)[:100],
value=partial
)
return [choice]
@remindme_group.command(
name=_p('cmd:remindme_in', "in"),
description=_p(
'cmd:remindme_in|desc',
"Schedule a reminder for a given amount of time in the future."
)
)
@appcmds.rename(
time=_p('cmd:remindme_in|param:time', "time"),
reminder=_p('cmd:remindme_in|param:reminder', "reminder"),
every=_p('cmd:remindme_in|param:every', "repeat_every"),
)
@appcmds.describe(
time=_p(
'cmd:remindme_in|param:time|desc',
"How far into the future to set the reminder (e.g. 1 day 10h 5m)."
),
reminder=_p(
'cmd:remindme_in|param:reminder|desc',
"What should the reminder be?"
),
every=_p(
'cmd:remindme_in|param:every|desc',
"How often to repeat this reminder. (e.g. 1 day, or 2h)"
)
)
async def cmd_remindme_in(
self,
ctx: LionContext,
time: Transform[int, DurationTransformer(60)],
reminder: appcmds.Range[str, 1, 2000],
every: Optional[Transform[int, DurationTransformer(60)]] = None
):
t = self.bot.translator.t
try:
remind_at = utc_now() + dt.timedelta(seconds=time)
reminder = await self.create_reminder(
userid=ctx.author.id,
remind_at=remind_at,
content=reminder,
message_link=ctx.message.jump_url,
interval=every,
)
embed = reminder.set_response
except UserInputError as e:
embed = discord.Embed(
title=t(_p(
'cmd:remindme_in|error|title',
"Could not create reminder!"
)),
description=e.msg,
colour=discord.Colour.brand_red()
)
await ctx.reply(
embed=embed,
ephemeral=True
)