From d549e5f5eb0efce7bacc713b1bb999ac3cba4440 Mon Sep 17 00:00:00 2001
From: Joshua Potter <jrpotter2112@gmail.com>
Date: Fri, 1 Dec 2023 07:10:58 -0700
Subject: [PATCH] Export blitz and bullet ratings.

---
 app/chesscom.py |  6 ++++++
 app/exporter.py | 24 ++++++++++++++++++++----
 app/lichess.py  | 28 +++++++++++++++++++++++-----
 3 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/app/chesscom.py b/app/chesscom.py
index 84b43cb..eba33da 100644
--- a/app/chesscom.py
+++ b/app/chesscom.py
@@ -187,3 +187,9 @@ class Exporter(BaseExporter):
 
     def export_rapid(self) -> Union[int, None]:
         return self.stats_json.get("rapid", {}).get("rating")
+
+    def export_blitz(self) -> Union[int, None]:
+        return self.stats_json.get("lightning", {}).get("rating")
+
+    def export_bullet(self) -> Union[int, None]:
+        return self.stats_json.get("bullet", {}).get("rating")
diff --git a/app/exporter.py b/app/exporter.py
index 3b3733f..f6df1e1 100644
--- a/app/exporter.py
+++ b/app/exporter.py
@@ -4,8 +4,18 @@ from typing_extensions import TypedDict
 
 
 class Export(TypedDict, total=False):
-    # The coach's rapid rating as listed on the site they were sourced from.
+    # The coach's rapid rating relative to the site they were sourced from.
     rapid: int
+    # The coach's blitz rating relative to the site they were sourced from.
+    blitz: int
+    # The coach's bullet rating relative to the site they were sourced from.
+    bullet: int
+
+
+def _insert(export: Export, key: str, value: any):
+    if value is None:
+        return
+    export[key] = value
 
 
 class BaseExporter(Repo):
@@ -16,13 +26,19 @@ class BaseExporter(Repo):
     def export_rapid(self) -> Union[int, None]:
         raise NotImplementedError()
 
+    def export_blitz(self) -> Union[int, None]:
+        raise NotImplementedError()
+
+    def export_bullet(self) -> Union[int, None]:
+        raise NotImplementedError()
+
     def export(self) -> Export:
         """Transform coach-specific data into uniform format."""
         export: Export = {}
 
-        rapid = self.export_rapid()
-        if rapid:
-            export["rapid"] = rapid
+        _insert(export, "rapid", self.export_rapid())
+        _insert(export, "blitz", self.export_blitz())
+        _insert(export, "bullet", self.export_bullet())
 
         self.log(
             [
diff --git a/app/lichess.py b/app/lichess.py
index 4d6dbf8..8bad4d7 100644
--- a/app/lichess.py
+++ b/app/lichess.py
@@ -6,7 +6,7 @@ import os.path
 from app.repo import AnsiColor, Site
 from app.scraper import BaseScraper
 from app.exporter import BaseExporter
-from bs4 import BeautifulSoup
+from bs4 import BeautifulSoup, SoupStrainer
 from typing import List
 
 
@@ -163,6 +163,12 @@ class Scraper(BaseScraper):
         return True
 
 
+def _stats_filter(elem, attrs):
+    """Includes only relevant segments of the `stats.html` file."""
+    if "sub-ratings" in attrs.get("class", ""):
+        return True
+
+
 class Exporter(BaseExporter):
     def __init__(self, username: str):
         super().__init__(site=Site.LICHESS.value, username=username)
@@ -170,18 +176,30 @@ class Exporter(BaseExporter):
         self.stats_soup = None
         try:
             with open(self.path_coach_file(username, "stats.html"), "r") as f:
-                self.stats_soup = BeautifulSoup(f.read(), "html.parser")
+                stats_strainer = SoupStrainer(_stats_filter)
+                self.stats_soup = BeautifulSoup(
+                    f.read(), "html.parser", parse_only=stats_strainer
+                )
         except FileNotFoundError:
             pass
 
     def export_rapid(self):
+        return self._find_rating("rapid")
+
+    def export_blitz(self):
+        return self._find_rating("blitz")
+
+    def export_bullet(self):
+        return self._find_rating("bullet")
+
+    def _find_rating(self, name):
         if self.stats_soup is None:
             return None
 
-        rapid = self.stats_soup.find("a", href=f"/@/{self.username}/perf/rapid")
-        if rapid is None:
+        anchor = self.stats_soup.find("a", href=f"/@/{self.username}/perf/{name}")
+        if anchor is None:
             return None
-        rating = rapid.find("rating")
+        rating = anchor.find("rating")
         if rating is None:
             return None
         strong = rating.find("strong")