Add election cog

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<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 = [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))