diff --git a/src/cogs/election.py b/src/cogs/election.py new file mode 100644 index 0000000..98437bc --- /dev/null +++ b/src/cogs/election.py @@ -0,0 +1,134 @@ +""" +This module is only meant to be loaded during election times. +""" +import asyncio +import logging +import re + +import discord +import httpx +from bs4 import BeautifulSoup +from discord.ext import commands + + +SPAN_REGEX = re.compile( + r"^(?P[a-zA-Z]+)\s(?P\d+)\scouncillors\s(?P\d+)\scouncillors" + r"\s(?P(gained|lost))$" +) +MULTI: dict[str, int] = { + "gained": 1, + "lost": -1 +} + + +class ElectionCog(commands.Cog): + SOURCE = "https://bbc.com/" + HEADERS = { + "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp," + "image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7", + "Accept-Language": "en-GB,en-US;q=0.9,en;q=0.8", + "Priority": "u=0, i", + "Sec-Ch-Ua": "\"Chromium\";v=\"124\", \"Google Chrome\";v=\"124\", \"Not-A.Brand\";v=\"99\"", + "Sec-Ch-Ua-Mobile": "?0", + "Sec-Ch-Ua-Platform": "\"Linux\"", + "Sec-Fetch-Dest": "document", + "Sec-Fetch-Mode": "navigate", + "Sec-Fetch-Site": "none", + "Sec-Fetch-User": "?1", + "Upgrade-Insecure-Requests": "1", + "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 " + "Safari/537.36", + } + + def __init__(self, bot): + self.bot = bot + self.log = logging.getLogger("jimmy.cogs.election") + + def process_soup(self, soup: BeautifulSoup) -> dict[str, list[int]] | None: + good_soup = soup.find(attrs={"data-testid": "election-banner-results-bar"}) + if not good_soup: + return + + css: str = [x for x in good_soup.find_all("style") if "Results" in x.get_text()][0].get_text() + + def find_colour(style_name: str, want: str = "background-colour") -> str | None: + index = css.index(style_name) + len(style_name) + 1 + value = "" + for char in css[index:]: + if char == "}": + break + value += char + attributes = filter(value.split(";")) + parsed = {} + for attr in attributes: + name, val = attr.split(":") + parsed[name] = val + if want in parsed: + return parsed[name] + + results: dict[str, list[int]] = {} + for child_ul in good_soup.children: + child_ul: BeautifulSoup + span = child_ul.find("span") + if not span: + self.log.warning("%r did not have a 'span' element.", child_ul) + continue + + text = span.get_text() + groups = SPAN_REGEX.match(text) + if groups: + groups = groups.groupdict() + else: + self.log.warning( + "Found span element (%r), however resolved text (%r) did not match regex.", + span, text + ) + continue + + results[str(groups["party"])] = [ + int(groups["councillors"]), + int(groups["net"]) * MULTI[groups["net_change"]], + int(find_colour(child_ul.next)[1:], base=16) + ] + return results + + @commands.slash_command(name="election") + async def get_election_results(self, ctx: discord.ApplicationContext): + """Gets the current election results""" + await ctx.defer() + async with httpx.AsyncClient(headers=self.HEADERS) as client: + response = await client.get(self.SOURCE) + if response.status_code != 200: + return await ctx.respond( + "Sorry, I can't do that right now (HTTP %d while fetching results from BBC)" % response.status_code + ) + + # noinspection PyTypeChecker + soup = await asyncio.to_thread(BeautifulSoup, response.text) + results = await self.bot.loop.run_in_executor(None, self.process_soup, soup) + if results: + date = ctx.message.created_at.date().strftime("%B %Y") + colour_scores = {} + embed = discord.Embed( + title="Election results - " + date, + url="https://bbc.co.uk/" + ) + embed.set_footer(text="Source from bbc.co.uk.") + description_parts = [] + + for party_name, values in results.items(): + councillors, net, colour = values + colour_scores[party_name] = councillors + description_parts.append( + f"**{party_name}**: {net:,}" + ) + + top_party = list(sorted(colour_scores.keys(), key=lambda k: colour_scores[k], reverse=True))[0] + embed.colour = discord.Colour(results[top_party][2]) + return await ctx.respond(embed=embed) + else: + return await ctx.respond("Unable to get election results at this time.") + + +def setup(bot): + bot.add_cog(ElectionCog(bot))