diff --git a/bot.py b/bot.py index ec66f41..ef937e0 100644 --- a/bot.py +++ b/bot.py @@ -2,17 +2,15 @@ import discord from discord.ext import tasks import os import sqlite3 +from collections import defaultdict from aiohttp import ClientSession import json client = discord.Client() conn = sqlite3.connect('data.db') - -def setup(): - conn.execute(''' - CREATE TABLE IF NOT EXISTS - guild(id int primary key, channel int, leaderboard int, last_sync text) - ''') +star_names = { '1': 'silver', '2': 'gold' } +aoc_session = os.getenv('AOC_SESSION') +bot_token = os.getenv('AOC_BOT_TOKEN') @client.event async def on_ready(): @@ -32,44 +30,120 @@ async def on_message(message): 'Updated {0.guild} to use channel {0.channel} and leaderboard {1}!' .format(message, leaderboard)) -def gather_stars(member): +def setup(): + ''' + Create all the necessary tables and all that. + ''' + + conn.execute(''' + CREATE TABLE IF NOT EXISTS + guild(id int primary key, channel int, leaderboard int, last_sync text) + ''') + +def member_stars(member): + ''' + Find the stars attributed to a single member + in the JSON response. Takes a JSON object + from the 'members' dict. + ''' + star_list = [] for day in member['completion_day_level']: stars = member['completion_day_level'][day] - if "1" in stars: - star_list.append((member['name'], int(day), 1, stars['1']['get_star_ts'])) - if "2" in stars: - star_list.append((member['name'], int(day), 2, stars["2"]['get_star_ts'])) + for (star, name) in star_names.items(): + if star not in stars: continue + star_list.append( + (member['name'], day, name, stars[star]['get_star_ts'])) return star_list -def perform_member_diff(old_member, new_member): - old_stars = gather_stars(old_member) - new_stars = gather_stars(new_member) - return [star for star in new_stars if star not in old_stars] +def all_stars(json): + ''' + Find all the stars in a given JSON response. + The stars are not ordered in any way. + ''' -def perform_diff(old_json, new_json): - if old_json is None: - old_json = {'members': {}} - else: - old_json = json.loads(old_json) - new_json = json.loads(new_json) - diff = { 'join': [], 'new_stars': [] } - for new_member in new_json['members'].keys(): - if new_member not in old_json['members']: - diff['join'].append(new_json['members'][new_member]['name']) - continue - member_diff = perform_member_diff(old_json['members'][new_member], new_json['members'][new_member]) - diff['new_stars'] += member_diff - return diff + return [star for m in json['members'].values() for star in member_stars(m)] + +def stars_by_day(stars): + ''' + Groups stars by day, then by type (silver/gold). + Helpful when trying to figure out if a given star + was "early" (first/second/third) for a particular puzzle. + ''' + + by_day = defaultdict(lambda: defaultdict(lambda: [])) + for star in stars: + by_day[star[1]][star[2]].append(star) + return by_day + +def read_json(s): + ''' + Try to parse a JSON response. If the argument + is None (which it can be if we're pulling from a fresh database) + fill it with just enough dummy enough to make the update + algorithms work. + ''' + + if s is None: return {'members': {}} + return json.loads(s) + +def detect_early_stars(stars): + ''' + Detects stars for each puzzle that are 'first', 'second', and 'third'. + This helps make the announcements a little more interesting! + ''' + + by_day = stars_by_day(stars) + early_stars = [] + for (day, names) in by_day.items(): + for (name, stars) in names.items(): + stars.sort(key=lambda star: int(star[3])) + early_stars += list(zip(['first', 'second', 'third'], stars)) + return early_stars + +def find_updates(old_json, new_json): + ''' + Find interesting differences between the old JSON and the new JSON, + both given as either strings or None. + + Interesting differences include member joins (for whom + we do not print new stars, since they weren't on the learderboard + when they won them), early stars (first/second/third) and other + stars (anyone winning a star at any point). + ''' + + old_json = read_json(old_json) + new_json = read_json(new_json) + + old_stars = all_stars(old_json) + new_stars = all_stars(new_json) + + join = [m for m in new_json['members'] if m not in old_json['members']] + early_stars = [ star for star in detect_early_stars(new_stars) + if star[1] not in old_stars + and star[1][0] not in join ] + skip_stars = [early_star for (place, early_star) in early_stars] + ann_stars = [ star for star in new_stars + if star not in skip_stars + and star not in old_stars + and star[0] not in join ] + return { 'join': join, 'early_stars': early_stars, 'ann_stars': ann_stars } + +async def send_updates(id, channel, diff): + ''' + Send updates via the Discord client given a diff + produced by `find_updates`. + ''' -async def send_diff(id, channel, diff): announcements = [] for user in diff['join']: full_text = "User {0} joined the leaderboard!".format(user) announcements.append(full_text) - for (user, day, star, ts) in diff['new_stars']: - star_name = "silver" if star == 1 else "gold" - full_text = "{0} won the {1} star from day {2}!".format(user, star_name, day) + for (place, (user, day, star, ts)) in diff['early_stars']: + full_text = "{0} won the {1} {2} star from day {3}!".format(user, place, star, day) + announcements.append(full_text) + for (user, day, star, ts) in diff['ann_stars']: + full_text = "{0} won the {1} star from day {2}!".format(user, star, day) announcements.append(full_text) if len(announcements) != 0: @@ -79,17 +153,17 @@ async def send_diff(id, channel, diff): @tasks.loop(minutes=1.0) async def update_aoc(): url_string = "https://adventofcode.com/2020/leaderboard/private/view/{0}.json" - cookies = { 'session': os.getenv("AOC_SESSION") } + cookies = { 'session': aoc_session } async with ClientSession(cookies=cookies) as session: for (id, channel, leaderboard, json) in conn.execute('select * from guild'): async with session.get(url_string.format(leaderboard)) as resp: new_json = await resp.text() - diff = perform_diff(json, new_json) - await send_diff(id, channel, diff) + diff = find_updates(json, new_json) + await send_updates(id, channel, diff) conn.execute('update guild set last_sync=? where id=?', (new_json, id)) conn.commit() setup() update_aoc.start() -client.run(os.environ['AOC_BOT_TOKEN']) +client.run(bot_token)