college-bot-v2/src/cogs/election.py

"""
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<party>[a-zA-Z]+)\s(?P<councillors>\d+)\scouncillors\s(?P<net>\d+)\scouncillors"
    r"\s(?P<net_change>(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 = "\n".join([x.get_text() for x in soup.find_all("style")])

        def find_colour(style_name: str, want: str = "background-colour") -> str | None:
            self.log.info("Looking for style %r", style_name)
            index = css.index(style_name) + len(style_name) + 1
            value = ""
            for char in css[index:]:
                if char == "}":
                    break
                value += char
            attributes = filter(None, value.split(";"))
            parsed = {}
            for attr in attributes:
                name, val = attr.split(":")
                parsed[name] = val
            if want in parsed:
                return parsed[want]

        results: dict[str, list[int]] = {}
        for child_ul in good_soup.children:
            child_ul: BeautifulSoup
            span = child_ul.find("span", recursive=False)
            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["class"][0])[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, follow_redirects=True)
            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, "html.parser")
            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))