feat(api): Add type validation for API Responses

src/api/pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
     "python-dotenv>=1.0.1",
     "pydantic-settings>=2.6.1",
     "eliot>=1.16.0",
+    "type-extensions>=0.1.2",
 ]
 
 [tool.ruff.lint]

src/api/src/create_book.py

@@ -1,4 +1,5 @@
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
+from typing_extensions import TypedDict
 import re
 import unicodedata
 import logging
@@ -11,7 +12,7 @@ from dotenv import load_dotenv
 from ebooklib import epub
 from ebooklib.epub import EpubBook
 from bs4 import BeautifulSoup
-from pydantic import model_validator, field_validator
+from pydantic import TypeAdapter, model_validator, field_validator
 from pydantic_settings import BaseSettings
 from aiohttp import ClientResponseError
 from aiohttp_client_cache.session import CachedSession
@@ -157,13 +158,48 @@ async def wp_get_cookies(username: str, password: str) -> dict:
                 return cookies
 
 
+# --- Models --- #
+
+
+class Language(TypedDict):
+    name: str
+
+
+class User(TypedDict):
+    username: str
+
+
+class Part(TypedDict):
+    id: int
+    title: str
+
+
+class Story(TypedDict):
+    id: str
+    title: str
+    createDate: str
+    modifyDate: str
+    language: Language
+    user: User
+    description: str
+    cover: str
+    completed: bool
+    tags: List[str]
+    mature: bool
+    url: str
+    parts: List[Part]
+    isPaywalled: bool
+
+
+story_ta = TypeAdapter(Story)
+
 # --- API Calls --- #
 
 
 @backoff.on_exception(backoff.expo, ClientResponseError, max_time=15)
 async def fetch_story_from_partId(
     part_id: int, cookies: Optional[dict] = None
-) -> Tuple[int, dict]:
+) -> Tuple[str, Story]:
     """Return a Story ID from a Part ID."""
     with start_action(action_type="api_fetch_storyFromPartId"):
         async with CachedSession(
@@ -176,11 +212,11 @@ async def fetch_story_from_partId(
 
                 body = await response.json()
 
-        return body["groupId"], body["group"]
+        return str(body["groupId"]), story_ta.validate_python(body["group"])
 
 
 @backoff.on_exception(backoff.expo, ClientResponseError, max_time=15)
-async def retrieve_story(story_id: int, cookies: Optional[dict] = None) -> dict:
+async def retrieve_story(story_id: int, cookies: Optional[dict] = None) -> Story:
     """Taking a story_id, return its information from the Wattpad API."""
     with start_action(action_type="api_fetch_story", story_id=story_id):
         async with CachedSession(
@@ -193,7 +229,7 @@ async def retrieve_story(story_id: int, cookies: Optional[dict] = None) -> dict:
 
                 body = await response.json()
 
-        return body
+        return story_ta.validate_python(body)
 
 
 @backoff.on_exception(backoff.expo, ClientResponseError, max_time=15)
@@ -231,7 +267,7 @@ async def fetch_cover(url: str) -> bytes:
 # --- EPUB Generation --- #
 
 
-def set_metadata(book: EpubBook, data: dict) -> None:
+def set_metadata(book: EpubBook, data: Story) -> None:
     """Set book metadata."""
     book.add_author(data["user"]["username"])
 
@@ -252,7 +288,7 @@ def set_metadata(book: EpubBook, data: dict) -> None:
     )
 
 
-async def set_cover(book: EpubBook, data: dict) -> None:
+async def set_cover(book: EpubBook, data: Story) -> None:
     """Set book cover."""
     book.set_cover("cover.jpg", await fetch_cover(data["cover"]))
     chapter = epub.EpubHtml(
@@ -263,7 +299,7 @@ async def set_cover(book: EpubBook, data: dict) -> None:
 
 async def add_chapters(
     book: EpubBook,
-    data: dict,
+    data: Story,
     download_images: bool = False,
     cookies: Optional[dict] = None,
 ):

src/api/src/main.py

@@ -72,7 +72,6 @@ app.add_middleware(RequestCancelledMiddleware)
 class DownloadMode(Enum):
     story = "story"
     part = "part"
-    collection = "collection"
 
 
 @app.get("/")
@@ -148,9 +147,7 @@ async def handle_download(
         logger.error(f"Retrieved story id ({story_id=})")
 
         book = epub.EpubBook()
-
         set_metadata(book, metadata)
-
         await set_cover(book, metadata)
 
         async for title in add_chapters(

src/api/uv.lock

@@ -195,6 +195,7 @@ dependencies = [
     { name = "pydantic-settings" },
     { name = "python-dotenv" },
     { name = "rich" },
+    { name = "type-extensions" },
 ]
 
 [package.metadata]
@@ -207,6 +208,7 @@ requires-dist = [
     { name = "pydantic-settings", specifier = ">=2.6.1" },
     { name = "python-dotenv", specifier = ">=1.0.1" },
     { name = "rich", specifier = ">=13.9.4" },
+    { name = "type-extensions", specifier = ">=0.1.2" },
 ]
 
 [[package]]
@@ -894,6 +896,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/96/00/2b325970b3060c7cecebab6d295afe763365822b1306a12eeab198f74323/starlette-0.41.3-py3-none-any.whl", hash = "sha256:44cedb2b7c77a9de33a8b74b2b90e9f50d11fcf25d8270ea525ad71a25374ff7", size = 73225 },
 ]
 
+[[package]]
+name = "type-extensions"
+version = "0.1.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/80/d0/18bf3a92af3f6a0fc38f874e858d707242ac5c14148ff1c4ee1615397e73/type_extensions-0.1.2.tar.gz", hash = "sha256:7c7b54eba6d7401ad5e69ecec6b7c767d7d7aae9b2b7e56249bc7bcaf833161f", size = 16102 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/43/dc/707d386bbd830d72fdd58c06f1d8a935d2553498f6868cd7de141addbe06/type_extensions-0.1.2-py2.py3-none-any.whl", hash = "sha256:fc9c8506f87b6227eeb43795061d61571b66a07df6fbcacf31a17d2cb2050153", size = 5566 },
+]
+
 [[package]]
 name = "typing-extensions"
 version = "4.12.2"