diff --git a/src/api/src/create_book.py b/src/api/src/create_book.py index 18386a5..6a240b8 100644 --- a/src/api/src/create_book.py +++ b/src/api/src/create_book.py @@ -7,6 +7,7 @@ import logging import tempfile import unicodedata from os import environ +from io import BytesIO from enum import Enum from base64 import b64encode import backoff @@ -515,14 +516,14 @@ class EPUBGenerator: # create spine self.epub.spine = ["nav"] + chapters - def dump(self) -> tempfile._TemporaryFileWrapper[bytes]: + def dump(self) -> BytesIO: # Thanks https://stackoverflow.com/a/75398222 - temp_file = tempfile.NamedTemporaryFile(suffix=".epub", delete=True) - epub.write_epub(temp_file, self.epub) + buffer = BytesIO() + epub.write_epub(buffer, self.epub) - temp_file.seek(0) + buffer.seek(0) - return temp_file + return buffer class PDFGenerator: @@ -634,6 +635,9 @@ style="margin-bottom: 1rem;">""".format( for image_container in html.find_all("p", {"data-media-type": "image"}): # Find all images, download them if download_images, else clear them (else wkhtmltopdf _might_ fetch them) img = image_container.findChild("img") + if not img: + image_container.decompose() # If empty, delete parent (ex:
) + continue source = img.get("src") if not download_images and source: img["src"] = "" @@ -744,10 +748,12 @@ style="margin-bottom: 1rem;">""".format( for chapter in chapters: chapter.file.close() - def dump(self) -> PDFGenerator: + def dump(self) -> BytesIO: self.file.seek(0) + buffer = BytesIO(self.file.read()) + self.file.close() - return self + return buffer # ------ # diff --git a/src/api/src/main.py b/src/api/src/main.py index ae2aa32..372af1b 100644 --- a/src/api/src/main.py +++ b/src/api/src/main.py @@ -3,7 +3,6 @@ from typing import Optional import asyncio from pathlib import Path -from io import BytesIO from enum import Enum from eliot import start_action from aiohttp import ClientResponseError @@ -192,12 +191,10 @@ async def handle_download( ): ... - book_file = book.dump().file - book_bytes = book_file.read() - book_file.close() # Deletes tempfile + book_buffer = book.dump() return StreamingResponse( - BytesIO(book_bytes), + book_buffer, media_type=media_type, headers={ "Content-Disposition": f'attachment; filename="{slugify(metadata["title"])}_{story_id}{"_images" if download_images else ""}.{format.value}"' # Thanks https://stackoverflow.com/a/72729058