feat(api): Add DownloadFormat type, restructure utils

2024-12-06 07:27:56 +00:00
parent 0f6cdd91a9
commit 0835992b23
2 changed files with 112 additions and 100 deletions
@@ -269,45 +269,49 @@ async def fetch_cover(url: str) -> bytes:
 # --- EPUB Generation --- #


-def set_metadata(book: EpubBook, data: Story) -> None:
-    """Set book metadata."""
-    book.add_author(data["user"]["username"])
+class EPUBGenerator:
+    def __init__(self, epub: EpubBook, data: Story):
+        self.epub = epub
+        self.data = data

-    book.add_metadata("DC", "title", data["title"])
-    book.add_metadata("DC", "description", data["description"])
-    book.add_metadata("DC", "date", data["createDate"])
-    book.add_metadata("DC", "modified", data["modifyDate"])
-    book.add_metadata("DC", "language", data["language"]["name"])
+        # set metadata
+        self.epub.add_author(data["user"]["username"])

-    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(
-        None, "meta", "", {"name": "completed", "content": str(int(data["completed"]))}
+        self.epub.add_metadata(
+            None,
+            "meta",
+            "",
+            {"name": "completed", "content": str(int(data["completed"]))},
        )

-
-async def set_cover(book: EpubBook, data: Story) -> None:
+    async def set_cover(self) -> 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,
-    data: Story,
+        self,
        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())
-    book.add_item(epub.EpubNav())
+        self.epub.add_item(epub.EpubNcx())
+        self.epub.add_item(epub.EpubNav())

        # create spine
-    book.spine = ["nav"] + chapters
+        self.epub.spine = ["nav"] + chapters
@@ -14,9 +14,7 @@ from fastapi.staticfiles import StaticFiles
 from ebooklib import epub
 from create_book import (
    retrieve_story,
-    set_cover,
-    set_metadata,
-    add_chapters,
+    EPUBGenerator,
    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()
-        set_metadata(book, metadata)
-        await set_cover(book, metadata)
+        match format:
+            case DownloadFormat.epub:
+                book = EPUBGenerator(epub.EpubBook(), metadata)
+                await book.set_cover()

-        async for title in add_chapters(
-            book, metadata, download_images=download_images, cookies=cookies
+                async for title in book.add_chapters(
+                    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")