Archived
generated from imcompiling/discord-bot-template
96 lines
3.6 KiB
Python
96 lines
3.6 KiB
Python
"""Contains commands for birthday feature."""
|
|
|
|
from datetime import datetime
|
|
|
|
import nextcord
|
|
from nextcord.ext import commands
|
|
from pytz import UTC
|
|
from sqlmodel import Session, col, select
|
|
|
|
from off_cooldown_bot.app import birthday_list
|
|
from off_cooldown_bot.database.db import engine
|
|
from off_cooldown_bot.database.models import UserBirthday
|
|
from off_cooldown_bot.utils import GUILDS
|
|
|
|
|
|
class BirthdayCog(commands.Cog, name="Birthday Commands"):
|
|
"""Class holds birthday commands."""
|
|
|
|
@nextcord.slash_command(guild_ids=GUILDS)
|
|
async def birthday(self, interaction: nextcord.Interaction) -> None:
|
|
"""Parent slash command for birthday. Never called."""
|
|
|
|
@birthday.subcommand(description="Set your birthday month and day for birthday fun!")
|
|
async def set(
|
|
self,
|
|
interaction: nextcord.Interaction,
|
|
month: int = nextcord.SlashOption(required=True, min_value=1, max_value=12, description="Your birth month."),
|
|
day: int = nextcord.SlashOption(required=True, min_value=1, max_value=31, description="Your birth day."),
|
|
) -> None:
|
|
"""
|
|
Set a user's birthday in the database. If birthday is today, adds to the birthday_list.
|
|
|
|
Args:
|
|
interaction (nextcord.Interaction): Interaction that invoked the command.
|
|
month (int, optional): The month a user was born. Is SlashOption.
|
|
day (int, optional): The day a user was born. Is SlashOption.
|
|
|
|
"""
|
|
# Ensure that a person used the command.
|
|
if interaction.user is None:
|
|
return
|
|
|
|
await interaction.response.defer(ephemeral=True)
|
|
|
|
# check if valid day
|
|
today = datetime.now(tz=UTC)
|
|
try:
|
|
birthday = datetime(today.year, month, day, tzinfo=UTC)
|
|
except ValueError:
|
|
await interaction.followup.send(f"Invalid date: {month}/{day}.", ephemeral=True)
|
|
return
|
|
|
|
# If the birthday is today, add to birthday_list
|
|
if today.date() == birthday.date():
|
|
birthday_list.append(interaction.user.id)
|
|
|
|
new_birthday = UserBirthday(user_id=interaction.user.id, month=month, day=day)
|
|
with Session(engine) as session:
|
|
stmt = select(UserBirthday).where(UserBirthday.user_id == interaction.user.id)
|
|
if (new_birthday := session.scalar(stmt)) is None:
|
|
new_birthday = UserBirthday(user_id=interaction.user.id, month=month, day=day)
|
|
session.add(new_birthday)
|
|
else:
|
|
new_birthday.month = month
|
|
new_birthday.day = day
|
|
session.commit()
|
|
|
|
await interaction.followup.send(f"Birthday set to {month}/{day}.", ephemeral=True)
|
|
|
|
@birthday.subcommand(description="Remove yourself from the birthday list.")
|
|
async def clear(self, interaction: nextcord.Interaction) -> None:
|
|
"""
|
|
Remove the user from the birthday database.
|
|
|
|
Args:
|
|
interaction (nextcord.Interaction): Interaction that invoked the command.
|
|
|
|
"""
|
|
# Ensure that a person used the command.
|
|
if interaction.user is None:
|
|
return
|
|
|
|
if interaction.user.id in birthday_list:
|
|
birthday_list.remove(interaction.user.id)
|
|
|
|
with Session(engine) as session:
|
|
stmt = select(UserBirthday).where(col(UserBirthday.user_id) == interaction.user.id)
|
|
if (user := session.scalar(stmt)) is None:
|
|
await interaction.response.send_message("No birthday to clear.", ephemeral=True)
|
|
return
|
|
|
|
session.delete(user)
|
|
session.commit()
|
|
|
|
await interaction.response.send_message("Birthday cleared.", ephemeral=True)
|