241 lines
8.3 KiB
Python
241 lines
8.3 KiB
Python
import discord
|
|
import pendulum
|
|
import asyncio
|
|
import traceback
|
|
import re
|
|
|
|
from discord.ext import commands
|
|
from . import utils
|
|
|
|
|
|
tzmap = {
|
|
'us-central': pendulum.timezone('US/Central'),
|
|
'eu-central': pendulum.timezone('Europe/Paris'),
|
|
'hongkong': pendulum.timezone('Hongkong'),
|
|
|
|
}
|
|
|
|
|
|
def sort_birthdays(bds):
|
|
# First sort the birthdays based on the comparison of the actual date
|
|
bds = sorted(bds, key=lambda x: x['birthday'])
|
|
# We want to split this into birthdays after and before todays date
|
|
# We can then use this to sort based on "whose is closest"
|
|
later_bds = []
|
|
previous_bds = []
|
|
# Loop through each birthday
|
|
for bd in bds:
|
|
# If it is after or equal to today, insert into our later list
|
|
if bd['birthday'] >= pendulum.today().date():
|
|
later_bds.append(bd)
|
|
# Otherwise, insert into our previous list
|
|
else:
|
|
previous_bds.append(bd)
|
|
# At this point we have 2 lists, in order, one from all of dates before today, and one after
|
|
# So all we need to do is put them in order all of "laters" then all of "befores"
|
|
return later_bds + previous_bds
|
|
|
|
|
|
def parse_string(date):
|
|
year = pendulum.now().year
|
|
month = None
|
|
day = None
|
|
month_map = {
|
|
"january": 1,
|
|
"jan": 1,
|
|
"february": 2,
|
|
"feb": 2,
|
|
"march": 3,
|
|
"mar": 3,
|
|
"april": 4,
|
|
"apr": 4,
|
|
"may": 5,
|
|
"june": 6,
|
|
"jun": 6,
|
|
"july": 7,
|
|
"jul": 7,
|
|
"august": 8,
|
|
"aug": 8,
|
|
"september": 9,
|
|
"sep": 9,
|
|
"october": 10,
|
|
"oct": 10,
|
|
"november": 11,
|
|
"nov": 11,
|
|
"december": 12,
|
|
"dec": 12,
|
|
}
|
|
|
|
num_re = re.compile("^(\d+)[a-z]*$")
|
|
|
|
for part in [x.lower() for x in date.split()]:
|
|
match = num_re.match(part)
|
|
if match:
|
|
day = int(match.group(1))
|
|
elif part in month_map:
|
|
month = month_map.get(part)
|
|
if month and day:
|
|
return pendulum.date(year, month, day)
|
|
|
|
|
|
class Birthday:
|
|
def __init__(self, bot):
|
|
self.bot = bot
|
|
self.task = self.bot.loop.create_task(self.birthday_task())
|
|
|
|
def get_birthdays_for_server(self, server, today=False):
|
|
bds = self.bot.db.load('birthdays')
|
|
# Get a list of the ID's to compare against
|
|
member_ids = [str(m.id) for m in server.members]
|
|
|
|
# Now create a list comparing to the server's list of member IDs
|
|
bds = [
|
|
bd
|
|
for member_id, bd in bds.items()
|
|
if str(member_id) in member_ids
|
|
]
|
|
|
|
_entries = []
|
|
|
|
for bd in bds:
|
|
if not bd['birthday']:
|
|
continue
|
|
|
|
day = parse_string(bd['birthday'])
|
|
# tz = tzmap.get(server.region)
|
|
# Check if it's today, and we want to only get todays birthdays
|
|
if (today and day.date() == pendulum.today().date()) or not today:
|
|
# If so, get the member and add them to the entry
|
|
member = server.get_member(int(bd['member_id']))
|
|
_entries.append({
|
|
'birthday': day,
|
|
'member': member
|
|
})
|
|
|
|
return sort_birthdays(_entries)
|
|
|
|
async def birthday_task(self):
|
|
while True:
|
|
try:
|
|
await self.notify_birthdays()
|
|
except Exception as error:
|
|
with open("error_log", 'a') as f:
|
|
traceback.print_tb(error.__traceback__, file=f)
|
|
print('{0.__class__.__name__}: {0}'.format(error), file=f)
|
|
finally:
|
|
# Every 12 hours, this is not something that needs to happen often
|
|
await asyncio.sleep(60 * 60 * 12)
|
|
|
|
async def notify_birthdays(self):
|
|
tfilter = {'birthdays_allowed': True}
|
|
servers = await self.bot.db.actual_load('server_settings', table_filter=tfilter)
|
|
for s in servers:
|
|
server = self.bot.get_guild(int(s['server_id']))
|
|
if not server:
|
|
continue
|
|
|
|
# Set our default to either the one set
|
|
default_channel_id = s.get('notifications', {}).get('default')
|
|
# If it is has been overriden by picarto notifications setting, use this
|
|
channel_id = s.get('notifications', {}).get('birthday') or default_channel_id
|
|
if not channel_id:
|
|
continue
|
|
|
|
# Now get the channel based on that ID
|
|
channel = server.get_channel(int(channel_id))
|
|
|
|
bds = self.get_birthdays_for_server(server, today=True)
|
|
for bd in bds:
|
|
try:
|
|
await channel.send("It is {}'s birthday today! "
|
|
"Wish them a happy birthday! \N{SHORTCAKE}".format(bd['member'].mention))
|
|
except (discord.Forbidden, discord.HTTPException, AttributeError):
|
|
pass
|
|
|
|
@commands.group(aliases=['birthdays'], invoke_without_command=True)
|
|
@commands.guild_only()
|
|
@utils.can_run(send_messages=True)
|
|
async def birthday(self, ctx, *, member: discord.Member = None):
|
|
"""A command used to view the birthdays on this server; or a specific member's birthday
|
|
|
|
EXAMPLE: !birthdays
|
|
RESULT: A printout of the birthdays from everyone on this server"""
|
|
if member:
|
|
date = self.bot.db.load('birthdays', key=member.id, pluck='birthday')
|
|
if date:
|
|
await ctx.send("{}'s birthday is {}".format(member.display_name, date))
|
|
else:
|
|
await ctx.send("I do not have {}'s birthday saved!".format(member.display_name))
|
|
else:
|
|
# Get this server's birthdays
|
|
bds = self.get_birthdays_for_server(ctx.message.guild)
|
|
# Create entries based on the user's display name and their birthday
|
|
entries = ["{} ({})".format(bd['member'].display_name, bd['birthday'].strftime("%B %-d")) for bd in bds]
|
|
# Create our pages object
|
|
try:
|
|
pages = utils.Pages(self.bot, message=ctx.message, entries=entries, per_page=5)
|
|
pages.title = "Birthdays for {}".format(ctx.message.guild.name)
|
|
await pages.paginate()
|
|
except utils.CannotPaginate as e:
|
|
await ctx.send(str(e))
|
|
|
|
@birthday.command(name='add')
|
|
@utils.can_run(send_messages=True)
|
|
async def _add_bday(self, ctx, *, date):
|
|
"""Used to link your birthday to your account
|
|
|
|
EXAMPLE: !birthday add December 1st
|
|
RESULT: I now know your birthday is December 1st"""
|
|
if len(date.split()) != 2:
|
|
await ctx.send("Please provide date in a valid format, such as December 1st!")
|
|
return
|
|
|
|
date = parse_string(date)
|
|
if date is None:
|
|
await ctx.send("Please provide date in a valid format, such as December 1st!")
|
|
return
|
|
|
|
date = date.strftime("%B %-d")
|
|
entry = {
|
|
'member_id': str(ctx.message.author.id),
|
|
'birthday': date
|
|
}
|
|
await self.bot.db.save('birthdays', entry)
|
|
await ctx.send("I have just saved your birthday as {}".format(date))
|
|
|
|
@birthday.command(name='remove')
|
|
@utils.can_run(send_messages=True)
|
|
async def _remove_bday(self, ctx):
|
|
"""Used to unlink your birthday to your account
|
|
|
|
EXAMPLE: !birthday remove
|
|
RESULT: I have magically forgotten your birthday"""
|
|
entry = {
|
|
'member_id': str(ctx.message.author.id),
|
|
'birthday': None
|
|
}
|
|
await self.bot.db.save('birthdays', entry)
|
|
await ctx.send("I don't know your birthday anymore :(")
|
|
|
|
@birthday.command(name='alerts', aliases=['notifications'])
|
|
@commands.guild_only()
|
|
@utils.can_run(manage_guild=True)
|
|
async def birthday_alerts_channel(self, ctx, channel: discord.TextChannel):
|
|
"""Sets the notifications channel for birthday notifications
|
|
|
|
EXAMPLE: !birthday alerts #birthday
|
|
RESULT: birthday notifications will go to this channel
|
|
"""
|
|
entry = {
|
|
'server_id': str(ctx.message.guild.id),
|
|
'notifications': {
|
|
'birthday': str(channel.id)
|
|
}
|
|
}
|
|
await self.bot.db.save('server_settings', entry)
|
|
await ctx.send("All birthday notifications will now go to {}".format(channel.mention))
|
|
|
|
|
|
def setup(bot):
|
|
bot.add_cog(Birthday(bot))
|