From d03ece6f58a5fef6d2c316b9a6e3f00b6926f9c5 Mon Sep 17 00:00:00 2001 From: stijndcl Date: Fri, 23 Sep 2022 14:25:13 +0200 Subject: [PATCH] Create reminders, fix bugs in schedule parsing --- .../versions/a64876b41af2_add_reminders.py | 38 ++++++++++++ database/crud/reminders.py | 42 ++++++++++++++ database/enums.py | 11 ++-- database/schemas.py | 16 +++++ didier/cogs/fun.py | 2 +- didier/cogs/meta.py | 27 +++++++++ didier/cogs/school.py | 5 +- didier/cogs/tasks.py | 58 ++++++++++++++++++- didier/data/embeds/schedules.py | 17 +++--- 9 files changed, 199 insertions(+), 17 deletions(-) create mode 100644 alembic/versions/a64876b41af2_add_reminders.py create mode 100644 database/crud/reminders.py diff --git a/alembic/versions/a64876b41af2_add_reminders.py b/alembic/versions/a64876b41af2_add_reminders.py new file mode 100644 index 0000000..40d48b5 --- /dev/null +++ b/alembic/versions/a64876b41af2_add_reminders.py @@ -0,0 +1,38 @@ +"""Add reminders + +Revision ID: a64876b41af2 +Revises: c1f9ee875616 +Create Date: 2022-09-23 13:37:10.331840 + +""" +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "a64876b41af2" +down_revision = "c1f9ee875616" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "reminders", + sa.Column("reminder_id", sa.Integer(), nullable=False), + sa.Column("user_id", sa.BigInteger(), nullable=True), + sa.Column("category", sa.Enum("LES", name="remindercategory"), nullable=False), + sa.ForeignKeyConstraint( + ["user_id"], + ["users.user_id"], + ), + sa.PrimaryKeyConstraint("reminder_id"), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("reminders") + # ### end Alembic commands ### diff --git a/database/crud/reminders.py b/database/crud/reminders.py new file mode 100644 index 0000000..007a779 --- /dev/null +++ b/database/crud/reminders.py @@ -0,0 +1,42 @@ +from typing import Optional + +from sqlalchemy import delete, select +from sqlalchemy.ext.asyncio import AsyncSession + +from database.crud.users import get_or_add_user +from database.enums import ReminderCategory +from database.schemas import Reminder + +__all__ = ["get_all_reminders_for_category", "toggle_reminder"] + + +async def get_all_reminders_for_category(session: AsyncSession, category: ReminderCategory) -> list[Reminder]: + """Get a list of all Reminders for a given category""" + statement = select(Reminder).where(Reminder.category == category) + return (await session.execute(statement)).scalars().all() + + +async def toggle_reminder(session: AsyncSession, user_id: int, category: ReminderCategory) -> bool: + """Switch a category on/off + + Returns the new value for the category + """ + await get_or_add_user(session, user_id) + + select_statement = select(Reminder).where(Reminder.user_id == user_id).where(Reminder.category == category) + reminder: Optional[Reminder] = (await session.execute(select_statement)).scalar_one_or_none() + + # No reminder set yet + if reminder is None: + reminder = Reminder(user_id=user_id, category=category) + session.add(reminder) + await session.commit() + + return True + + # Reminder found -> delete it + delete_statement = delete(Reminder).where(Reminder.reminder_id == reminder.reminder_id) + await session.execute(delete_statement) + await session.commit() + + return False diff --git a/database/enums.py b/database/enums.py index e2cf565..3eec739 100644 --- a/database/enums.py +++ b/database/enums.py @@ -1,11 +1,14 @@ import enum -__all__ = ["TaskType"] +__all__ = ["ReminderCategory", "TaskType"] + + +class ReminderCategory(enum.IntEnum): + """Enum for reminder categories""" + + LES = enum.auto() -# There is a bug in typeshed that causes an incorrect PyCharm warning -# https://github.com/python/typeshed/issues/8286 -# noinspection PyArgumentList class TaskType(enum.IntEnum): """Enum for the different types of tasks""" diff --git a/database/schemas.py b/database/schemas.py index f497b2d..3fd8ce4 100644 --- a/database/schemas.py +++ b/database/schemas.py @@ -37,6 +37,7 @@ __all__ = [ "Link", "MemeTemplate", "NightlyData", + "Reminder", "Task", "UforaAnnouncement", "UforaCourse", @@ -219,6 +220,18 @@ class NightlyData(Base): user: User = relationship("User", back_populates="nightly_data", uselist=False, lazy="selectin") +class Reminder(Base): + """Something that a user should be reminded of""" + + __tablename__ = "reminders" + + reminder_id: int = Column(Integer, primary_key=True) + user_id: int = Column(BigInteger, ForeignKey("users.user_id")) + category: enums.ReminderCategory = Column(Enum(enums.ReminderCategory), nullable=False) + + user: User = relationship("User", back_populates="reminders", uselist=False, lazy="selectin") + + class Task(Base): """A Didier task""" @@ -303,6 +316,9 @@ class User(Base): nightly_data: NightlyData = relationship( "NightlyData", back_populates="user", uselist=False, lazy="selectin", cascade="all, delete-orphan" ) + reminders: list[Reminder] = relationship( + "Reminder", back_populates="user", uselist=True, lazy="selectin", cascade="all, delete-orphan" + ) wordle_guesses: list[WordleGuess] = relationship( "WordleGuess", back_populates="user", uselist=True, lazy="selectin", cascade="all, delete-orphan" ) diff --git a/didier/cogs/fun.py b/didier/cogs/fun.py index fe78275..25d0baf 100644 --- a/didier/cogs/fun.py +++ b/didier/cogs/fun.py @@ -83,7 +83,7 @@ class Fun(commands.Cog): @memes_slash.command(name="generate") async def memegen_slash(self, interaction: discord.Interaction, template: str): - """Generate a meme with template `template`.""" + """Generate a meme.""" async with self.client.postgres_session as session: result = expect(await get_meme_by_name(session, template), entity_type="meme", argument=template) diff --git a/didier/cogs/meta.py b/didier/cogs/meta.py index 7b7a732..ae4ea81 100644 --- a/didier/cogs/meta.py +++ b/didier/cogs/meta.py @@ -4,6 +4,8 @@ from typing import Optional from discord.ext import commands +from database.crud.reminders import toggle_reminder +from database.enums import ReminderCategory from didier import Didier @@ -20,6 +22,31 @@ class Meta(commands.Cog): """Get Didier's latency.""" return await ctx.reply(f"Polo! {round(self.client.latency * 1000)}ms", mention_author=False) + @commands.command(name="remind", aliases=["remindme"]) + async def remind(self, ctx: commands.Context, category: str): + """Make Didier send you reminders every day.""" + category = category.lower() + + available_categories = [ + ( + "les", + ReminderCategory.LES, + ) + ] + + for name, category_mapping in available_categories: + if name == category: + async with self.client.postgres_session as session: + new_value = await toggle_reminder(session, ctx.author.id, category_mapping) + + toggle = "on" if new_value else "off" + return await ctx.reply( + f"Reminders for category `{name}` have been toggled {toggle}.", mention_author=False + ) + + # No match found + return await ctx.reply(f"`{category}` is not a supported category.", mention_author=False) + @commands.command(name="source", aliases=["src"]) async def source(self, ctx: commands.Context, *, command_name: Optional[str] = None): """Get a link to the source code of Didier. diff --git a/didier/cogs/school.py b/didier/cogs/school.py index aacdc4d..a9b6b7a 100644 --- a/didier/cogs/school.py +++ b/didier/cogs/school.py @@ -11,7 +11,7 @@ from didier import Didier from didier.data.apis.hydra import fetch_menu from didier.data.embeds.deadlines import Deadlines from didier.data.embeds.hydra import no_menu_found -from didier.data.embeds.schedules import Schedule, get_schedule_for_user +from didier.data.embeds.schedules import Schedule, get_schedule_for_day from didier.exceptions import HTTPException, NotInMainGuildException from didier.utils.discord.converters.time import DateTransformer from didier.utils.discord.flags.school import StudyGuideFlags @@ -55,10 +55,11 @@ class School(commands.Cog): try: member_instance = to_main_guild_member(self.client, ctx.author) + roles = {role.id for role in member_instance.roles} # Always make sure there is at least one schedule in case it returns None # this allows proper error messages - schedule = get_schedule_for_user(self.client, member_instance, day_dt) or Schedule() + schedule = (get_schedule_for_day(self.client, day_dt) or Schedule()).personalize(roles) return await ctx.reply(embed=schedule.to_embed(day=day_dt), mention_author=False) diff --git a/didier/cogs/tasks.py b/didier/cogs/tasks.py index 37b1786..1657be6 100644 --- a/didier/cogs/tasks.py +++ b/didier/cogs/tasks.py @@ -9,10 +9,16 @@ from overrides import overrides import settings from database import enums from database.crud.birthdays import get_birthdays_on_day +from database.crud.reminders import get_all_reminders_for_category from database.crud.ufora_announcements import remove_old_announcements from database.crud.wordle import set_daily_word +from database.schemas import Reminder from didier import Didier -from didier.data.embeds.schedules import Schedule, parse_schedule_from_content +from didier.data.embeds.schedules import ( + Schedule, + get_schedule_for_day, + parse_schedule_from_content, +) from didier.data.embeds.ufora.announcements import fetch_ufora_announcements from didier.decorators.tasks import timed_task from didier.utils.discord.checks import is_owner @@ -44,6 +50,7 @@ class Tasks(commands.Cog): self._tasks = { "birthdays": self.check_birthdays, "schedules": self.pull_schedules, + "reminders": self.reminders, "ufora": self.pull_ufora_announcements, "remove_ufora": self.remove_old_ufora_announcements, "wordle": self.reset_wordle_word, @@ -61,6 +68,7 @@ class Tasks(commands.Cog): self.remove_old_ufora_announcements.start() # Start other tasks + self.reminders.start() self.reset_wordle_word.start() self.pull_schedules.start() @@ -135,7 +143,7 @@ class Tasks(commands.Cog): async with self.client.postgres_session as session: for data in settings.SCHEDULE_DATA: if data.schedule_url is None: - return + continue async with self.client.http_session.get(data.schedule_url) as response: # If a schedule couldn't be fetched, log it and move on @@ -180,6 +188,51 @@ class Tasks(commands.Cog): async def _before_ufora_announcements(self): await self.client.wait_until_ready() + async def _send_les_reminders(self, entries: list[Reminder]): + today = datetime.date(year=2022, month=9, day=26) + + # Create the main schedule for the day once here, to avoid doing it repeatedly + daily_schedule = get_schedule_for_day(self.client, today) + + # No class today + if not daily_schedule: + return + + for entry in entries: + member = self.client.main_guild.get_member(entry.user_id) + if not member: + continue + + roles = {role.id for role in member.roles} + personal_schedule = daily_schedule.personalize(roles) + + # No class today + if not personal_schedule: + continue + + await member.send(embed=personal_schedule.to_embed(day=today)) + + # @tasks.loop(time=SOCIALLY_ACCEPTABLE_TIME) + @tasks.loop(hours=3) + async def reminders(self, **kwargs): + """Send daily reminders to people""" + _ = kwargs + + async with self.client.postgres_session as session: + for category in enums.ReminderCategory: + entries = await get_all_reminders_for_category(session, category) + if not entries: + continue + + # This is slightly ugly, but it's the best way to go about it + # There won't be a lot of categories anyway + if category == enums.ReminderCategory: + await self._send_les_reminders(entries) + + @reminders.before_loop + async def _before_reminders(self): + await self.client.wait_until_ready() + @tasks.loop(hours=24) async def remove_old_ufora_announcements(self): """Remove all announcements that are over 1 week old, once per day""" @@ -200,6 +253,7 @@ class Tasks(commands.Cog): @check_birthdays.error @pull_schedules.error @pull_ufora_announcements.error + @reminders.error @remove_old_ufora_announcements.error @reset_wordle_word.error async def _on_tasks_error(self, error: BaseException): diff --git a/didier/data/embeds/schedules.py b/didier/data/embeds/schedules.py index ed03a33..39e16de 100644 --- a/didier/data/embeds/schedules.py +++ b/didier/data/embeds/schedules.py @@ -22,7 +22,7 @@ from didier.utils.types.datetime import LOCAL_TIMEZONE, int_to_weekday, time_str from didier.utils.types.string import leading from settings import ScheduleType -__all__ = ["Schedule", "get_schedule_for_user", "parse_schedule_from_content", "parse_schedule"] +__all__ = ["Schedule", "get_schedule_for_day", "parse_schedule_from_content", "parse_schedule"] @dataclass @@ -48,6 +48,10 @@ class Schedule(EmbedBaseModel): def personalize(self, roles: set[int]) -> Schedule: """Personalize a schedule for a user, only adding courses they follow""" + # If the schedule is already empty, just return instantly + if not self.slots: + return Schedule() + personal_slots = set() for slot in self.slots: role_found = slot.role_id is not None and slot.role_id in roles @@ -104,10 +108,9 @@ class ScheduleSlot: def __post_init__(self): """Fix some properties to display more nicely""" # Re-format the location data - room, building, campus = re.search(r"(.*)\. Gebouw (.*)\. Campus (.*)\. ", self.location).groups() + room, building, campus = re.search(r"(.*)\. (?:Gebouw )?(.*)\. (?:Campus )?(.*)\. ", self.location).groups() room = room.replace("PC / laptoplokaal ", "PC-lokaal") self.location = f"{campus} {building} {room}" - self._hash = hash(f"{self.course.course_id} {str(self.start_time)}") @property @@ -132,14 +135,12 @@ class ScheduleSlot: return self._hash == other._hash -def get_schedule_for_user(client: Didier, member: discord.Member, day_dt: date) -> Optional[Schedule]: - """Get a user's schedule""" - roles: set[int] = {role.id for role in member.roles} - +def get_schedule_for_day(client: Didier, day_dt: date) -> Optional[Schedule]: + """Get a schedule for an entire day""" main_schedule: Optional[Schedule] = None for schedule in client.schedules.values(): - personalized_schedule = schedule.on_day(day_dt).personalize(roles) + personalized_schedule = schedule.on_day(day_dt) if not personalized_schedule: continue