feat(api): Add DownloadFormat type, restructure utils

src/api/src/create_book.py

@@ -269,45 +269,49 @@ async def fetch_cover(url: str) -> bytes:
 # --- EPUB Generation --- #
 
 
-def set_metadata(book: EpubBook, data: Story) -> None:
+class EPUBGenerator:
-    """Set book metadata."""
+    def __init__(self, epub: EpubBook, data: Story):
-    book.add_author(data["user"]["username"])
+        self.epub = epub
+        self.data = data
 
-    book.add_metadata("DC", "title", data["title"])
+        # set metadata
-    book.add_metadata("DC", "description", data["description"])
+        self.epub.add_author(data["user"]["username"])
-    book.add_metadata("DC", "date", data["createDate"])
-    book.add_metadata("DC", "modified", data["modifyDate"])
-    book.add_metadata("DC", "language", data["language"]["name"])
 
-    book.add_metadata(
+        self.epub.add_metadata("DC", "title", data["title"])
+        self.epub.add_metadata("DC", "description", data["description"])
+        self.epub.add_metadata("DC", "date", data["createDate"])
+        self.epub.add_metadata("DC", "modified", data["modifyDate"])
+        self.epub.add_metadata("DC", "language", data["language"]["name"])
+
+        self.epub.add_metadata(
             None, "meta", "", {"name": "tags", "content": ", ".join(data["tags"])}
         )
-    book.add_metadata(
+        self.epub.add_metadata(
             None, "meta", "", {"name": "mature", "content": str(int(data["mature"]))}
         )
-    book.add_metadata(
+        self.epub.add_metadata(
-        None, "meta", "", {"name": "completed", "content": str(int(data["completed"]))}
+            None,
+            "meta",
+            "",
+            {"name": "completed", "content": str(int(data["completed"]))},
         )
 
-
+    async def set_cover(self) -> None:
-async def set_cover(book: EpubBook, data: Story) -> None:
         """Set book cover."""
-    book.set_cover("cover.jpg", await fetch_cover(data["cover"]))
+        self.epub.set_cover("cover.jpg", await fetch_cover(self.data["cover"]))
         chapter = epub.EpubHtml(
             file_name="titlepage.xhtml",  # Standard for cover page
         )
         chapter.set_content('<img src="cover.jpg">')
 
-
     async def add_chapters(
-    book: EpubBook,
+        self,
-    data: Story,
         download_images: bool = False,
         cookies: Optional[dict] = None,
     ):
         chapters = []
 
-    for cidx, part in enumerate(data["parts"]):
+        for cidx, part in enumerate(self.data["parts"]):
             content = await fetch_part_content(part["id"], cookies=cookies)
             title = part["title"]
 
@@ -315,7 +319,7 @@ async def add_chapters(
             chapter = epub.EpubHtml(
                 title=title,
                 file_name=f"{cidx}.xhtml",  # Used to be clean_title.xhtml, but that broke Arabic support as slugify turns arabic strings into '', leading to multiple files with the same name, breaking those chapters.
-            lang=data["language"]["name"],
+                lang=self.data["language"]["name"],
             )
 
             if download_images:
@@ -335,7 +339,7 @@ async def add_chapters(
                                 content=await response.read(),
                                 file_name=f"static/{cidx}/{idx}.jpeg",
                             )
-                        book.add_item(img)
+                            self.epub.add_item(img)
                             # Fetch image and pack
 
                             content = content.replace(
@@ -349,13 +353,13 @@ async def add_chapters(
             yield title  # Yield the chapter's title upon insertion preceeded by retrieval.
 
         for chapter in chapters:
-        book.add_item(chapter)
+            self.epub.add_item(chapter)
 
-    book.toc = chapters
+        self.epub.toc = chapters
 
         # Thanks https://github.com/aerkalov/ebooklib/blob/master/samples/09_create_image/create.py
-    book.add_item(epub.EpubNcx())
+        self.epub.add_item(epub.EpubNcx())
-    book.add_item(epub.EpubNav())
+        self.epub.add_item(epub.EpubNav())
 
         # create spine
-    book.spine = ["nav"] + chapters
+        self.epub.spine = ["nav"] + chapters

src/api/src/main.py

@@ -14,9 +14,7 @@ from fastapi.staticfiles import StaticFiles
 from ebooklib import epub
 from create_book import (
     retrieve_story,
-    set_cover,
+    EPUBGenerator,
-    set_metadata,
-    add_chapters,
     slugify,
     wp_get_cookies,
     fetch_story_from_partId,
@@ -69,6 +67,11 @@ class RequestCancelledMiddleware:
 app.add_middleware(RequestCancelledMiddleware)
 
 
+class DownloadFormat(Enum):
+    pdf = "pdf"
+    epub = "epub"
+
+
 class DownloadMode(Enum):
     story = "story"
     part = "part"
@@ -106,6 +109,7 @@ async def handle_download(
     download_id: int,
     download_images: bool = False,
     mode: DownloadMode = DownloadMode.story,
+    format: DownloadFormat = DownloadFormat.epub,
     username: Optional[str] = None,
     password: Optional[str] = None,
 ):
@@ -146,12 +150,13 @@ async def handle_download(
 
         logger.info(f"Retrieved story id ({story_id=})")
 
-        book = epub.EpubBook()
+        match format:
-        set_metadata(book, metadata)
+            case DownloadFormat.epub:
-        await set_cover(book, metadata)
+                book = EPUBGenerator(epub.EpubBook(), metadata)
+                await book.set_cover()
 
-        async for title in add_chapters(
+                async for title in book.add_chapters(
-            book, metadata, download_images=download_images, cookies=cookies
+                    download_images=download_images, cookies=cookies
                 ):
                     ...
 
@@ -161,7 +166,7 @@ async def handle_download(
                 )  # Thanks https://stackoverflow.com/a/75398222
 
                 # create epub file
-        epub.write_epub(temp_file, book, {})
+                epub.write_epub(temp_file, book.epub, {})
 
                 temp_file.file.seek(0)
                 book_data = temp_file.file.read()
@@ -174,6 +179,9 @@ async def handle_download(
                     },
                 )
 
+            case DownloadFormat.pdf:
+                ...
+
 
 app.mount("/", StaticFiles(directory=BUILD_PATH), "static")