fix(api): Use BytesIO when dumping generated book

This commit is contained in:
TheOnlyWayUp
2024-12-10 11:12:43 +00:00
parent 16c5a9216f
commit 016ad6209a
2 changed files with 15 additions and 12 deletions
+13 -7
View File
@@ -7,6 +7,7 @@ import logging
import tempfile import tempfile
import unicodedata import unicodedata
from os import environ from os import environ
from io import BytesIO
from enum import Enum from enum import Enum
from base64 import b64encode from base64 import b64encode
import backoff import backoff
@@ -515,14 +516,14 @@ class EPUBGenerator:
# create spine # create spine
self.epub.spine = ["nav"] + chapters self.epub.spine = ["nav"] + chapters
def dump(self) -> tempfile._TemporaryFileWrapper[bytes]: def dump(self) -> BytesIO:
# Thanks https://stackoverflow.com/a/75398222 # Thanks https://stackoverflow.com/a/75398222
temp_file = tempfile.NamedTemporaryFile(suffix=".epub", delete=True) buffer = BytesIO()
epub.write_epub(temp_file, self.epub) epub.write_epub(buffer, self.epub)
temp_file.seek(0) buffer.seek(0)
return temp_file return buffer
class PDFGenerator: class PDFGenerator:
@@ -634,6 +635,9 @@ style="margin-bottom: 1rem;">""".format(
for image_container in html.find_all("p", {"data-media-type": "image"}): 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) # Find all images, download them if download_images, else clear them (else wkhtmltopdf _might_ fetch them)
img = image_container.findChild("img") img = image_container.findChild("img")
if not img:
image_container.decompose() # If empty, delete parent (ex: <p data-image-layout="one-horizontal" data-media-type="image" data-p-id="bb6e18f2bb7d13f317bb6ccded04899b">        </p>)
continue
source = img.get("src") source = img.get("src")
if not download_images and source: if not download_images and source:
img["src"] = "" img["src"] = ""
@@ -744,10 +748,12 @@ style="margin-bottom: 1rem;">""".format(
for chapter in chapters: for chapter in chapters:
chapter.file.close() chapter.file.close()
def dump(self) -> PDFGenerator: def dump(self) -> BytesIO:
self.file.seek(0) self.file.seek(0)
buffer = BytesIO(self.file.read())
self.file.close()
return self return buffer
# ------ # # ------ #
+2 -5
View File
@@ -3,7 +3,6 @@
from typing import Optional from typing import Optional
import asyncio import asyncio
from pathlib import Path from pathlib import Path
from io import BytesIO
from enum import Enum from enum import Enum
from eliot import start_action from eliot import start_action
from aiohttp import ClientResponseError from aiohttp import ClientResponseError
@@ -192,12 +191,10 @@ async def handle_download(
): ):
... ...
book_file = book.dump().file book_buffer = book.dump()
book_bytes = book_file.read()
book_file.close() # Deletes tempfile
return StreamingResponse( return StreamingResponse(
BytesIO(book_bytes), book_buffer,
media_type=media_type, media_type=media_type,
headers={ headers={
"Content-Disposition": f'attachment; filename="{slugify(metadata["title"])}_{story_id}{"_images" if download_images else ""}.{format.value}"' # Thanks https://stackoverflow.com/a/72729058 "Content-Disposition": f'attachment; filename="{slugify(metadata["title"])}_{story_id}{"_images" if download_images else ""}.{format.value}"' # Thanks https://stackoverflow.com/a/72729058