From be9d15c566e74f5362c089be0ff684ccfd0101c5 Mon Sep 17 00:00:00 2001
From: Jon Betts <jon.betts@hypothes.is>
Date: Wed, 2 Aug 2023 18:33:25 +0100
Subject: [PATCH] Add id parsing to YouTube caption track objects

---
 .../via/services/youtube_api/models_test.py   | 38 +++++++++----------
 via/services/youtube_api/models.py            | 25 ++++++++++++
 2 files changed, 44 insertions(+), 19 deletions(-)

diff --git a/tests/unit/via/services/youtube_api/models_test.py b/tests/unit/via/services/youtube_api/models_test.py
index e257d034..11384bde 100644
--- a/tests/unit/via/services/youtube_api/models_test.py
+++ b/tests/unit/via/services/youtube_api/models_test.py
@@ -30,29 +30,29 @@ def test_from_v1_json(self, kind):
             }
         )
 
-    @pytest.mark.parametrize(
-        "caption_track,id_string",
+    CAPTION_TRACK_IDS = (
+        (CaptionTrack(language_code="en"), "en"),
+        (CaptionTrack(language_code="en", kind="asr"), "en.a"),
+        (CaptionTrack(language_code="en", name="Hello"), "en..SGVsbG8="),
+        (CaptionTrack(language_code="en", translated_language_code="fr"), "en...fr"),
+        # This combination isn't actually possible, but let's try everything at
+        # once
         (
-            (CaptionTrack(language_code="en"), "en"),
-            (CaptionTrack(language_code="en", kind="asr"), "en.a"),
-            (CaptionTrack(language_code="en", name="Hello"), "en..SGVsbG8="),
-            (
-                CaptionTrack(language_code="en", translated_language_code="fr"),
-                "en...fr",
-            ),
-            # This combination isn't actually possible, but let's try everything at
-            # once
-            (
-                CaptionTrack(
-                    language_code="en-gb",
-                    kind="asr",
-                    name="Name",
-                    translated_language_code="fr",
-                ),
-                "en-gb.a.TmFtZQ==.fr",
+            CaptionTrack(
+                language_code="en-gb",
+                kind="asr",
+                name="Name",
+                translated_language_code="fr",
             ),
+            "en-gb.a.TmFtZQ==.fr",
         ),
     )
+
+    @pytest.mark.parametrize("caption_track,id_string", CAPTION_TRACK_IDS)
+    def test_from_id(self, caption_track, id_string):
+        assert CaptionTrack.from_id(id_string) == caption_track
+
+    @pytest.mark.parametrize("caption_track,id_string", CAPTION_TRACK_IDS)
     def test_id(self, caption_track, id_string):
         assert caption_track.id == id_string
 
diff --git a/via/services/youtube_api/models.py b/via/services/youtube_api/models.py
index 5c24ddb3..71106487 100644
--- a/via/services/youtube_api/models.py
+++ b/via/services/youtube_api/models.py
@@ -1,6 +1,7 @@
 import base64
 from copy import deepcopy
 from dataclasses import dataclass, field
+from itertools import zip_longest
 from operator import attrgetter
 from typing import List, Optional
 
@@ -43,6 +44,30 @@ def from_v1_json(cls, data: dict):
             base_url=data["baseUrl"],
         )
 
+    @classmethod
+    def from_id(cls, id_string: str):
+        """Create a partially filled out track from and id string."""
+
+        data = dict(
+            zip_longest(
+                [
+                    "language_code",
+                    "auto_generated",
+                    "name",
+                    "translated_language_code",
+                ],
+                [part or None for part in id_string.split(".")],
+            )
+        )
+
+        if name := data.get("name"):
+            data["name"] = base64.b64decode(name.encode("utf-8")).decode("utf-8")
+
+        if data.pop("auto_generated", None):
+            data["kind"] = "asr"
+
+        return cls(**data)
+
     @property
     def id(self) -> str:  # pylint: disable=invalid-name
         if self.name: