|
| 1 | +import datetime |
| 2 | + |
| 3 | +import disnake |
| 4 | +import datetime as dt |
| 5 | +from disnake.ext import commands, tasks |
| 6 | +import random |
| 7 | + |
| 8 | +from utils.DBhandlers import BirthdayHandler |
| 9 | +from utils.exceptions import UserAlreadyExists, UserNotFound |
| 10 | +from utils.shortcuts import QuickEmb, sucEmb, errorEmb |
| 11 | +from utils.CONSTANTS import months, congrats_messages |
| 12 | +from utils.bot import OGIROID |
| 13 | + |
| 14 | + |
| 15 | +class Birthday(commands.Cog): |
| 16 | + def __init__(self, bot: OGIROID): |
| 17 | + self.bot = bot |
| 18 | + self.birthday: BirthdayHandler = None |
| 19 | + self.birthday_check.start() |
| 20 | + |
| 21 | + @commands.Cog.listener() |
| 22 | + async def on_ready(self): |
| 23 | + if not self.bot.ready_: |
| 24 | + self.birthday: BirthdayHandler = BirthdayHandler(self.bot, self.bot.db) |
| 25 | + |
| 26 | + def cog_unload(self): |
| 27 | + self.birthday_check.cancel() |
| 28 | + |
| 29 | + @commands.slash_command(name="birthday") |
| 30 | + async def birthday(self, inter: disnake.ApplicationCommandInteraction): |
| 31 | + pass |
| 32 | + |
| 33 | + @birthday.sub_command(name="set", description="Set your birthday. Cant be removed without Staff.") |
| 34 | + async def set( |
| 35 | + self, |
| 36 | + inter, |
| 37 | + day: int = commands.Param(name="day", ge=1, le=31, description="The day of your birthday. Select carefully."), |
| 38 | + month: str = commands.Param( |
| 39 | + name="month", |
| 40 | + description="The month of your birthday. Select carefully.", |
| 41 | + choices=months, |
| 42 | + ), |
| 43 | + ): |
| 44 | + if month is None or day is None: |
| 45 | + return await errorEmb(inter, "You need to provide a month and a day") |
| 46 | + if day < 1 or day > 31: |
| 47 | + return await errorEmb(inter, "The day must be between 1 and 31") |
| 48 | + |
| 49 | + birth_date = f"{day}/{month}" |
| 50 | + try: |
| 51 | + await self.birthday.create_user(inter.author.id, birth_date) |
| 52 | + except UserAlreadyExists: |
| 53 | + return await errorEmb(inter, "You already have a birthday set") |
| 54 | + |
| 55 | + await sucEmb(inter, f"Your birthday has been set to {birth_date}") |
| 56 | + |
| 57 | + @commands.has_permissions(manage_roles=True) |
| 58 | + @birthday.sub_command(name="edit", description="Edit a users birthday. Can only be done by Staff.") |
| 59 | + async def edit( |
| 60 | + self, |
| 61 | + inter, |
| 62 | + day: int = commands.Param(name="day", ge=1, le=31, description="The day of the birthday."), |
| 63 | + month: str = commands.Param( |
| 64 | + name="month", |
| 65 | + description="The month of the birthday.", |
| 66 | + choices=months, |
| 67 | + ), |
| 68 | + user: disnake.User = commands.Param(name="user", description="User to edit the birthday of."), |
| 69 | + ): |
| 70 | + try: |
| 71 | + await self.birthday.update_user(user.id, f"{day}/{month}") |
| 72 | + return await sucEmb(inter, f"Birthday has been updated to {day}/{month}") |
| 73 | + except UserNotFound: |
| 74 | + return await errorEmb(inter, "The User doesn't have a birthday set") |
| 75 | + |
| 76 | + @commands.has_permissions(manage_roles=True) |
| 77 | + @birthday.sub_command(name="remove", description="Remove a birthday. Can only be done by Staff.") |
| 78 | + async def remove( |
| 79 | + self, |
| 80 | + inter: disnake.ApplicationCommandInteraction, |
| 81 | + user: disnake.User = commands.Param(name="user", description="Removes the birthday of this user"), |
| 82 | + ): |
| 83 | + try: |
| 84 | + await self.birthday.delete_user(user.id) |
| 85 | + except UserNotFound: |
| 86 | + return await errorEmb(inter, "This user doesn't have a birthday set") |
| 87 | + |
| 88 | + await sucEmb(inter, "The birthday has been removed") |
| 89 | + |
| 90 | + @birthday.sub_command(name="get", description="Get the birthday of a user") |
| 91 | + async def get(self, inter, user: disnake.User = commands.Param(name="user", default=None)): |
| 92 | + if user is None: |
| 93 | + user = inter.author |
| 94 | + else: |
| 95 | + user = await self.bot.fetch_user(user.id) |
| 96 | + |
| 97 | + birthday = await self.birthday.get_user(user.id) |
| 98 | + if birthday is None: |
| 99 | + return await errorEmb(inter, "This user doesn't have a birthday set") |
| 100 | + |
| 101 | + next_birthday = datetime.datetime.strptime(birthday.birthday + f"/{dt.datetime.now().year}", "%d/%m/%Y") |
| 102 | + if next_birthday < datetime.datetime.now(): |
| 103 | + next_birthday = datetime.datetime.strptime(birthday.birthday + f"/{dt.datetime.now().year + 1}", "%d/%m/%Y") |
| 104 | + await QuickEmb( |
| 105 | + inter, |
| 106 | + f"{user.mention}'s birthday is in {(next_birthday - datetime.datetime.now()).days} Days." |
| 107 | + f" <t:{str(next_birthday.timestamp()).split('.')[0]}:D>", |
| 108 | + ).send() |
| 109 | + |
| 110 | + # @tasks.loop(time=[dt.time(dt.datetime.utcnow().hour, dt.datetime.utcnow().minute, dt.datetime.utcnow().second + 10)]) |
| 111 | + # ^ use this when testing birthdays |
| 112 | + @tasks.loop(time=[dt.time(8, 0, 0)]) |
| 113 | + # loops every day at 8:00 UTC time |
| 114 | + async def birthday_check(self): |
| 115 | + channel = self.bot.get_channel(self.bot.config.channels.birthdays) |
| 116 | + guild = self.bot.get_guild(self.bot.config.guilds.main_guild) |
| 117 | + if channel is None: |
| 118 | + return |
| 119 | + today = dt.datetime.utcnow().strftime("%d/%m") |
| 120 | + # Gets all users from the db |
| 121 | + users = await self.birthday.get_users() |
| 122 | + for user in users: |
| 123 | + member = await guild.fetch_member(user.user_id) |
| 124 | + # if the member is None, the user is not in the server anymore |
| 125 | + if member is None: |
| 126 | + continue |
| 127 | + |
| 128 | + # if the birthday is today, congratulate the user |
| 129 | + if user.birthday == today: |
| 130 | + await member.add_roles(guild.get_role(self.bot.config.roles.birthday)) |
| 131 | + congrats_msg = await channel.send(f"{random.choice(congrats_messages)} {member.mention}! 🎂") |
| 132 | + await congrats_msg.add_reaction("🥳") |
| 133 | + # If the birthday isn't today and the user still has the birthday role, remove it |
| 134 | + elif user.birthday != today and member.get_role(self.bot.config.roles.birthday) is not None: |
| 135 | + await member.remove_roles(guild.get_role(self.bot.config.roles.birthday)) |
| 136 | + |
| 137 | + |
| 138 | +def setup(bot): |
| 139 | + bot.add_cog(Birthday(bot)) |
0 commit comments