refactor(api): Rewrite PDFGenerator
This commit is contained in:
TheOnlyWayUp
@@ -0,0 +1,204 @@
|
||||
from base64 import b64encode
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from tempfile import NamedTemporaryFile, _TemporaryFileWrapper
|
||||
from typing import List, cast
|
||||
|
||||
from bs4 import BeautifulSoup, Tag
|
||||
from exiftool import ExifTool
|
||||
from jinja2 import Template
|
||||
from weasyprint import CSS, HTML
|
||||
from weasyprint.text.fonts import FontConfiguration
|
||||
|
||||
from ..models import Story
|
||||
from .types import AbstractGenerator
|
||||
|
||||
DATA_PATH = Path(__file__).parent / "pdf"
|
||||
ASSET_PATH = DATA_PATH / "assets"
|
||||
|
||||
COPYRIGHT_DATA = {
|
||||
1: {
|
||||
"name": "All Rights Reserved",
|
||||
"statement": "©️ {published_year} by {username}. All Rights Reserved.",
|
||||
"freedoms": "No reuse, redistribution, or modification without permission.",
|
||||
"printing": "Not allowed without explicit permission.",
|
||||
"asset": None,
|
||||
},
|
||||
2: {
|
||||
"name": "Public Domain",
|
||||
"statement": "This work is in the public domain. Originally published in {published_year} by {username}.",
|
||||
"freedoms": "Free to use for any purpose without permission.",
|
||||
"printing": "Allowed for personal or commercial purposes.",
|
||||
"asset": ASSET_PATH / "cc-zero.png",
|
||||
},
|
||||
3: {
|
||||
"name": "Creative Commons Attribution (CC-BY)",
|
||||
"statement": "©️ {published_year} by {username}. This work is licensed under a Creative Commons Attribution 4.0 International License.",
|
||||
"freedoms": "Allows reuse, redistribution, and modification with credit to the author.",
|
||||
"printing": "Allowed with proper credit.",
|
||||
"asset": ASSET_PATH / "by.png",
|
||||
},
|
||||
4: {
|
||||
"name": "CC Attribution NonCommercial (CC-BY-NC)",
|
||||
"statement": "©️ {published_year} by {username}. This work is licensed under a Creative Commons Attribution-NonCommercial 4.0 International License.",
|
||||
"freedoms": "Allows reuse and modification for non-commercial purposes with credit.",
|
||||
"printing": "Allowed for non-commercial purposes with proper credit.",
|
||||
"asset": ASSET_PATH / "by-nc.png",
|
||||
},
|
||||
5: {
|
||||
"name": "CC Attribution NonCommercial NoDerivs (CC-BY-NC-ND)",
|
||||
"statement": "©️ {published_year} by {username}. This work is licensed under a Creative Commons Attribution-NonCommercial-NoDerivs 4.0 International License.",
|
||||
"freedoms": "Allows sharing in original form for non-commercial purposes with credit; no modifications allowed.",
|
||||
"printing": "Allowed for non-commercial purposes in original form with proper credit.",
|
||||
"asset": ASSET_PATH / "by-nc-nd.png",
|
||||
},
|
||||
6: {
|
||||
"name": "CC Attribution NonCommercial ShareAlike (CC-BY-NC-SA)",
|
||||
"statement": "©️ {published_year} by {username}. This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License.",
|
||||
"freedoms": "Allows reuse and modification for non-commercial purposes under the same license, with credit.",
|
||||
"printing": "Allowed for non-commercial purposes with proper credit under the same license.",
|
||||
"asset": ASSET_PATH / "by-nc-sa.png",
|
||||
},
|
||||
7: {
|
||||
"name": "CC Attribution ShareAlike (CC-BY-SA)",
|
||||
"statement": "©️ {published_year} by {username}. This work is licensed under a Creative Commons Attribution-ShareAlike 4.0 International License.",
|
||||
"freedoms": "Allows reuse and modification for any purpose under the same license, with credit.",
|
||||
"printing": "Allowed with proper credit under the same license.",
|
||||
"asset": ASSET_PATH / "by-sa.png",
|
||||
},
|
||||
8: {
|
||||
"name": "CC Attribution NoDerivs (CC-BY-ND)",
|
||||
"statement": "©️ {published_year} by {username}. This work is licensed under a Creative Commons Attribution-NoDerivs 4.0 International License.",
|
||||
"freedoms": "Allows sharing in original form for any purpose with credit; no modifications allowed.",
|
||||
"printing": "Allowed in original form with proper credit.",
|
||||
"asset": ASSET_PATH / "by-nd.png",
|
||||
},
|
||||
} # Maps Wattpad Copyright IDs to their corresponding data.
|
||||
|
||||
with open(DATA_PATH / "stylesheet.css") as reader:
|
||||
STYLESHEET = reader.read()
|
||||
|
||||
|
||||
with open(DATA_PATH / "book.html") as reader:
|
||||
TEMPLATE = reader.read()
|
||||
|
||||
|
||||
class PDFGenerator(AbstractGenerator):
|
||||
def __init__(
|
||||
self,
|
||||
metadata: Story,
|
||||
part_trees: List[BeautifulSoup],
|
||||
cover: bytes,
|
||||
images: List[List[bytes]] | None,
|
||||
author: bytes,
|
||||
):
|
||||
self.story = metadata
|
||||
self.parts = part_trees
|
||||
self.cover = cover
|
||||
self.images = images
|
||||
self.author = author
|
||||
|
||||
self.book: _TemporaryFileWrapper = NamedTemporaryFile(suffix=".pdf")
|
||||
self.content = TEMPLATE
|
||||
|
||||
def generate_chapters(self) -> dict[int, str]:
|
||||
"""Return a dictionary of part_ids to content trees, with image URLs replaced with base64 encoded images if provided during initialization."""
|
||||
data: dict[int, str] = {}
|
||||
for idx, (part, tree) in enumerate(zip(self.story["parts"], self.parts)):
|
||||
if self.images:
|
||||
for img_idx, (img_data, img_tag) in enumerate(
|
||||
zip(self.images[idx], tree.find_all("img"))
|
||||
):
|
||||
img_tag[
|
||||
"src"
|
||||
] = f"data:image/jpg;base64,{b64encode(img_data).decode()}"
|
||||
|
||||
data[part["id"]] = tree.prettify()
|
||||
|
||||
return data
|
||||
|
||||
def populate_template(self, parts: dict[int, str]):
|
||||
"""Populate HTML Template with Story data."""
|
||||
copyright = COPYRIGHT_DATA[self.story["copyright"]]
|
||||
data = {
|
||||
"statement": copyright["statement"].format(
|
||||
username=self.story["user"]["username"],
|
||||
published_year=self.story["createDate"].split("-", 2)[0],
|
||||
),
|
||||
"author": self.story["user"]["username"],
|
||||
"freedoms": copyright["freedoms"],
|
||||
"printing": copyright["printing"],
|
||||
"book_id": self.story["id"],
|
||||
"book_title": self.story["title"],
|
||||
"cover": f"data:image/jpg;base64,{b64encode(self.cover).decode()}",
|
||||
"username": self.story["user"]["username"],
|
||||
"description": self.story["description"],
|
||||
"avatar": b64encode(self.author).decode(),
|
||||
"copyright": {
|
||||
"data": b64encode(copyright["asset"].read_bytes()).decode()
|
||||
if copyright["asset"]
|
||||
else "",
|
||||
"name": copyright["name"],
|
||||
},
|
||||
"parts": parts,
|
||||
}
|
||||
|
||||
self.content: str = Template(self.content).render(data)
|
||||
|
||||
def generate_pdf(self):
|
||||
"""Generate and write the PDF to a temporary file (self.book)."""
|
||||
font_config = FontConfiguration()
|
||||
|
||||
stylesheet_obj = CSS(string=STYLESHEET, font_config=font_config)
|
||||
|
||||
html_obj = HTML(string=self.content)
|
||||
html_obj.write_pdf(
|
||||
self.book.name, stylesheets=[stylesheet_obj], font_config=font_config
|
||||
)
|
||||
|
||||
def add_metadata(self):
|
||||
"""Write metadata to generated PDF file at self.book, using ExifTool."""
|
||||
|
||||
clean_description = (
|
||||
self.story["description"].strip().replace("\n", "$/")
|
||||
) # exiftool doesn't parse \ns correctly, they support $/ for the same instead. `
` is another option.
|
||||
|
||||
metadata = {
|
||||
"Author": self.story["user"]["username"],
|
||||
"Title": self.story["title"],
|
||||
"Subject": clean_description,
|
||||
"CreationDate": self.story["createDate"],
|
||||
"ModDate": self.story["modifyDate"],
|
||||
"Keywords": ",".join(self.story["tags"]),
|
||||
"Language": self.story["language"]["name"],
|
||||
"Completed": self.story["completed"],
|
||||
"MatureContent": self.story["mature"],
|
||||
"Producer": "Dhanush Rambhatla (TheOnlyWayUp - https://rambhat.la) and WattpadDownloader",
|
||||
} # As per https://exiftool.org/TagNames/PDF.html
|
||||
|
||||
with ExifTool(config_file=DATA_PATH / "exiftool.config") as et:
|
||||
# Custom configuration adds Completed and MatureContent tags.
|
||||
# exiftool logger logs executed command
|
||||
et.execute(
|
||||
*(
|
||||
[f"-{key}={value}" for key, value in metadata.items()]
|
||||
+ [
|
||||
"-overwrite_original",
|
||||
self.book.file.name,
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
def compile(self):
|
||||
parts = self.generate_chapters()
|
||||
self.populate_template(parts)
|
||||
self.generate_pdf()
|
||||
self.add_metadata()
|
||||
return True
|
||||
|
||||
def dump(self) -> BytesIO:
|
||||
self.book.seek(0)
|
||||
buffer = BytesIO(self.book.read())
|
||||
self.book.close()
|
||||
|
||||
return buffer
|
||||
Reference in New Issue
Block a user