diff --git a/src/api/pyproject.toml b/src/api/pyproject.toml index 6bdcc2e..b68d3ee 100644 --- a/src/api/pyproject.toml +++ b/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] diff --git a/src/api/src/create_book.py b/src/api/src/create_book.py index 4b01951..949d00b 100644 --- a/src/api/src/create_book.py +++ b/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, ): diff --git a/src/api/src/main.py b/src/api/src/main.py index 28534f9..b8d095f 100644 --- a/src/api/src/main.py +++ b/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( diff --git a/src/api/uv.lock b/src/api/uv.lock index 6785913..805275f 100644 --- a/src/api/uv.lock +++ b/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"