diff --git a/.github/workflows/python-tests.yaml b/.github/workflows/python-tests.yaml index 0e435a96..c3572405 100644 --- a/.github/workflows/python-tests.yaml +++ b/.github/workflows/python-tests.yaml @@ -12,7 +12,7 @@ name: Python tests # # The suites are separate because their dependencies are: docdb-studio and # check-tools each own a pyproject.toml + uv.lock, while ProcessKotlinWebsiteJSON -# runs against the root requirements.txt. +# and scripts/test_optimize_pdfs.py run against the root requirements.txt. # # check-tools is NOT run here, deliberately. Its suite covers the health checker # behind the daily docdb-regression-test.yaml cron and genuinely ought to be @@ -60,14 +60,16 @@ jobs: # it because no Python binding exposes a custom dictionary. # zstd: train_dictionary uses `zstd --train-fastcover` to build the # shared dictionary the migration tests need. - sudo apt-get install -y pngquant brotli zstd + # jbig2dec: pikepdf decodes JBIG2 through it, and optimize_pdfs.py's + # syntax check decodes every stream; its JBIG2Globals test needs it. + sudo apt-get install -y pngquant brotli zstd jbig2dec - name: Set up uv uses: astral-sh/setup-uv@v5 # scripts/test_migrate_tooltip_schema.py needs only the standard library, pytest and the brotli CLI (above). # scripts/test_trim_kotlin_json.py needs only brotli and Pillow (in requirements.txt) and the brotli CLI. - - name: 'ProcessKotlinWebsiteJSON + sync_kotlin_stdlib_docs + migrate_tooltip_schema + trim_kotlin_json + Android and Java loaders' + - name: 'ProcessKotlinWebsiteJSON + sync_kotlin_stdlib_docs + migrate_tooltip_schema + trim_kotlin_json + optimize_pdfs + Android and Java loaders' run: | uv run --with-requirements requirements.txt --with pytest -- \ python -m pytest \ @@ -75,6 +77,7 @@ jobs: scripts/sync_kotlin_stdlib_docs \ scripts/test_migrate_tooltip_schema.py \ scripts/test_trim_kotlin_json.py \ + scripts/test_optimize_pdfs.py \ ProcessDocs/AndroidDocs/AndroidDocsToJson \ scripts/sync_java_docs \ -q diff --git a/CLAUDE.md b/CLAUDE.md index 6d1db17b..f9510fd9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -251,7 +251,12 @@ shape: every tool here reads and writes that, and only the shipped copy holds th re-encodes the raster images SVG rows embed as data: URIs as WebP, image by image, by a quality rule; run on the shipped copy before the final dictionary re-mint; `compare_svg_renders.py` renders each changed SVG before and after in headless Chrome, through `` and ``, as - its visual check), and four steps that run in this order before the `--before` snapshot + its visual check), `optimize_pdfs.py` (ADFA-6629: rewrites the PDF rows - and, with `--file`, the Bookshelf + plugin's `.pdf.br` books - losslessly so the Brotli around them compresses them better: streams other + than images stored plain, Flate-wrapped JPEGs stored bare, Flate images PNG-predicted; every rewrite + checked object by object against the original; run on the shipped copy after `svg_images_to_webp.py`, + before the final dictionary re-mint), + and four steps that run in this order before the `--before` snapshot `update_media_references.py` diffs against: `remove_skin_font_comment.py` (cuts Help+Manual's commented-out Google Fonts note out of the `i/` pages until the skin fix, ADFA-6609, ships; first - its docstring says why), `bundle_remote_images.py` (stores the 19 diagrams API reference diff --git a/requirements.txt b/requirements.txt index 9d7c018c..59882a36 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,3 +8,4 @@ tqdm-loggable>=0.1.0 scour cairosvg markdown-it-py>=2.0 +pikepdf>=10 diff --git a/scripts/optimize_pdfs.py b/scripts/optimize_pdfs.py new file mode 100644 index 00000000..a5e433b0 --- /dev/null +++ b/scripts/optimize_pdfs.py @@ -0,0 +1,720 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.9" +# dependencies = [ +# "brotli", +# "Pillow", +# "pikepdf>=10", +# ] +# /// +""" +optimize_pdfs.py + +Rewrites PDFs losslessly so the Brotli they ship in compresses them better: +the PDF rows of a documentation database, or one PDF file such as a Bookshelf +plugin book (ADFA-6629). + +Why: every shipped PDF is already wrapped in Brotli (a Content row, or the +plugin's .pdf.br asset), but inside it most streams are Flate-compressed, and +Brotli cannot shrink Flate output. Storing those streams plain lets the outer +Brotli compress them instead. Classic PDF optimization (qpdf's object streams, +Flate recompressed at level 9) saved 0.1% here, and qpdf's --recompress-flate +strips the PNG predictors from images (Java Notes grew 60 KB), so it is not +used. + +The rewrite, per PDF: +- Every stream that is not an image and whose filters are all generalized + (Flate, LZW, ASCII85, ASCIIHex, RunLength) is stored plain: content streams, + fonts, ICC profiles, forms. +- A Flate image of 8-bit gray, 8-bit RGB or 1-bit samples is re-encoded as + PNG-predicted Flate (Pillow's optimized PNG encoder) when that is smaller, + and only when Pillow gives back exactly the same sample bytes (it zeroes the + pad bits at the end of a 1-bit row, which a reader may not). Image data + stays compressed, so the plain PDF stays reasonably small. +- A JPEG wrapped in Flate ([/FlateDecode /DCTDecode]) is stored as the bare + JPEG, its DCT decode parameters kept, when its Flate data is whole: the + checksum right and nothing after the end. The outer Brotli compresses these + JPEGs better than Flate does. +- Other images (bare JPEG, JPEG 2000, JBIG2, CCITT) are left as they are. +- Saved with object streams, no stream recompressed, and a deterministic /ID, + so the same input always gives the same bytes and a second run changes + nothing. + +Measured on the 2026-10-19 ship-time database and the Bookshelf books of +addons 6d14aa7, with pikepdf 10.5.1 (qpdf 12.3.2), Pillow 12.2.0 and brotli +1.2.0. Sizes as shipped: a database row against its dictionary, a book as +plain Brotli q11 w24. + + PDF before after plain PDF before -> after + d/PluginDevelopmentGuide.pdf 376,166 241,111 406,791 -> 1,124,176 + d/KeystoreGeneratorPluginDocumentation 279,750 192,150 900,857 -> 2,270,799 + d/SnippetsPluginDocumentation 193,215 129,435 304,385 -> 770,663 + d/BeepyPluginDocumentation 147,590 101,932 209,603 -> 601,784 + d/MarkdownPreviewerPluginDocumentation 97,427 76,757 114,266 -> 194,994 + d/APKViewerPluginDocumentation 93,606 79,950 112,701 -> 197,122 + AndroidNotesForProfessionals 8,810,306 6,320,438 9,797,404 -> 33,930,413 + JavaNotesForProfessionals 4,772,016 2,812,106 6,867,048 -> 25,480,261 + JavaJavaJava 4,162,347 2,258,547 4,752,514 -> 14,733,459 + BeejGuideToCProgramming 1,346,539 705,655 1,434,906 -> 6,771,870 + KotlinNotesForProfessionals 787,651 590,175 1,097,854 -> 3,718,498 + ModernCplusplusTutorialOuChangkun 388,053 278,975 417,807 -> 1,234,691 + PebbleTemplateGuide 90,299 50,466 154,606 -> 321,412 + +The six database rows: 1,187,754 -> 821,335 stored bytes. Packed as the +pipeline packs (rebuild_database.py --pack Content,TooltipTexts --page-size +4096) and set against the same database packed the same way without this +step, the file goes from 85,708,800 to 85,340,160 bytes (-368,640) and its .br +(brotli -q 11 -w 24) from 67,288,387 to 66,914,581 (-373,806). The nine +placeholder PDF rows (the Bookshelf stand-ins, under 2.4 KB) stay as they are: +their rewrite is larger. The seven books: 20,357,211 -> 13,016,362 bytes. +Every page of all thirteen (3,816) renders the same in poppler at 72 dpi +before and after, and pdfinfo -dests lists the same destinations. + +Before anything is written, the rewritten PDF must pass every check, or it is +left as it is and the report says why: +- walked from the trailer's /Root and /Info, both object graphs have the same + shape, and every dictionary is unchanged apart from a stream's /Filter, + /DecodeParms and /Length; +- every stream's decoded data is identical: fully decoded for generalized + filters (so a re-encoded image gives back the same samples), and for a JPEG, + JPEG 2000, JBIG2 or CCITT stream the payload that filter reads, with its + decode parameters (a stream among them, such as a JBIG2Globals, compared by + its data the same way); +- pikepdf's syntax check (qpdf's, which also decodes every stream, JPEGs and + JBIG2 included) finds nothing; +- the page count, the named destinations (the catalog's /Dests and the + /Names /Dests tree) and the outline, nested items, actions and levels + included, are unchanged; a stream in any of them is compared by its data. +A rewrite that is not smaller once compressed is left alone too. A stream qpdf +cannot decode without a warning is not rewritten, an image included: storing +the part qpdf could read would quietly change what a reader gets. Such a +stream, left as it is, usually fails the syntax check, and its PDF then stays +as it is. + +Left alone by policy, which is not a failure: an encrypted PDF, and a signed +one (its /AcroForm sets SigFlags bit 1, or it has a /Perms entry), because a +signature covers the file's bytes and every check would pass while it broke. + +pikepdf decodes JBIG2 through the jbig2dec program; where it is missing, the +syntax check cannot run on a PDF with a JBIG2 image, which then stays as it is +(refused, "it cannot be rewritten"). + +The qpdf command line's --check is stricter about name trees than the syntax +check: on the rewritten Beej and JavaJavaJava books it reports "skipping over +invalid kid", for empty dictionaries their /Dests name trees already hold in +the original, at the same places. The graph check proves those trees are +unchanged. + +The plain PDF grows (Android Notes from 9.8 MB to 33 MB), and Code On The Go +decodes a whole row into memory before serving it, so the phone test measures +the load. A size cap, if one is needed, belongs in keep_rewrite(). + +Database mode: every application/pdf Content row (chunk rows joined as +WebServer.kt joins them), decoded with the database's own dictionary, +rewritten, and re-encoded with that dictionary. Rows are updated in place: +the base row is never deleted and re-inserted, because the AddBook and +DeleteBook triggers would replace its curated Bookshelf entry. Continuation +rows are trimmed or added as the new size needs. A PDF row that is not +Brotli-compressed is skipped and named. Runs in one transaction (rolled back on +any error), backs the database up first, and VACUUMs afterwards. --dry-run does +all the work, the checks included, and writes nothing. The report ends with a +summary: how many PDFs were rewritten, and how many were left as they are +because they were not smaller, refused by a check (by check) or skipped by +policy (by reason). Exit 1 means nothing was written: a row could not be read, +or a check refused every PDF that policy did not skip. + +Pipeline position: on the copy that ships, after svg_images_to_webp.py and +before the final dictionary re-mint and rebuild_database.py --pack. It touches +only PDF rows. + +File mode: a .pdf, or a .pdf.br holding plain Brotli, to a .pdf.br of plain +Brotli at q11 w24 (no dictionary: the Bookshelf plugin inserts its asset bytes +as they are). When the rewrite is not smaller or policy skips the PDF, the +output holds the input PDF unchanged; when a check refuses it, nothing is +written and the exit is 1. The target keeps its mode; a new one gets +0666 less the umask, as open() would give it. + + uv run scripts/optimize_pdfs.py documentation.db --dry-run + uv run scripts/optimize_pdfs.py documentation.db [--workers N] [--verbose] + uv run scripts/optimize_pdfs.py --file books/JavaJavaJava.pdf.br out/JavaJavaJava.pdf.br +""" +from __future__ import annotations + +import argparse +import io +import os +import re +import sqlite3 +import stat +import struct +import sys +import tempfile +import zlib +from collections import Counter +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path + +import brotli +import pikepdf +from PIL import Image + +from optimize_db_media import BrotliCodec, backup_database, human, is_continuation_path, load_dictionary, reassemble +from update_media_references import replace_row + +PDF_TYPE = "application/pdf" +GENERALIZED = frozenset({"/FlateDecode", "/LZWDecode", "/ASCII85Decode", "/ASCIIHexDecode", "/RunLengthDecode"}) +# Stream dictionary keys a rewrite may change; every other key must stay as it was. +REWRITABLE_KEYS = frozenset({"/Filter", "/DecodeParms", "/Length"}) +# (components, bits per component) -> Pillow mode, for the images re-encoded as PNG-predicted Flate. +PNG_MODES = {(1, 8): "L", (3, 8): "RGB", (1, 1): "1"} +# The PNG predictor that lets each row choose its own filter, as a PNG encoder does. +PNG_OPTIMUM = 15 +# File mode's Brotli: the settings the Bookshelf books and the database download ship with. +FILE_QUALITY, FILE_LGWIN = 11, 24 + + +class LeftAsIs(Exception): + """Why a PDF is left as it is. `kind` names the check or the policy, for the summary. Both go in args, so the + exception survives the process pool's pickling.""" + + def __init__(self, message: str, kind: str): + super().__init__(message, kind) + + @property + def kind(self) -> str: + return self.args[1] + + def __str__(self) -> str: + return self.args[0] + + +class Refused(LeftAsIs): + """A rewritten PDF failed a check: a failure, which the summary and the exit code report.""" + + +class Skipped(LeftAsIs): + """A PDF this step does not rewrite by policy (encrypted, signed, not Brotli): not a failure.""" + + +def filters(stream: pikepdf.Object) -> list: + f = stream.get("/Filter") + if f is None: + return [] + return [str(x) for x in f] if isinstance(f, pikepdf.Array) else [str(f)] + + +def decode_parms(stream: pikepdf.Object, count: int) -> list: + """The stream's /DecodeParms as one entry per filter (None where absent).""" + p = stream.get("/DecodeParms") + if isinstance(p, pikepdf.Array): + return list(p) + [None] * (count - len(p)) + return [p] + [None] * (count - 1) if count else [] + + +def has_predictor(parms) -> bool: + return isinstance(parms, pikepdf.Dictionary) and int(parms.get("/Predictor", 1)) > 1 + + +def inflate(raw: bytes) -> bytes: + """Flate data decoded as a tolerant reader decodes it (a missing end of stream is not an error).""" + return zlib.decompressobj().decompress(raw) + + +def inflate_whole(raw: bytes) -> bytes | None: + """Flate data decoded only when it is whole: its checksum right and nothing after its end. Otherwise None, and + the stream is not rewritten.""" + d = zlib.decompressobj() + try: + data = d.decompress(raw) + d.flush() + except zlib.error: + return None + return data if d.eof and not d.unused_data else None + + +def components(colour_space) -> int | None: + if colour_space is None: + return None + name = str(colour_space[0] if isinstance(colour_space, pikepdf.Array) else colour_space) + if name in ("/DeviceGray", "/CalGray", "/Indexed", "/Separation"): + return 1 + if name in ("/DeviceRGB", "/CalRGB", "/Lab"): + return 3 + if name == "/ICCBased": + return int(colour_space[1].get("/N", 0)) + return None + + +def png_flate(stream: pikepdf.Object, samples: bytes) -> tuple | None: + """(PNG-predicted Flate data, its /DecodeParms) for an image stream's decoded samples, or None when the image is + not one this step re-encodes.""" + mask = bool(stream.get("/ImageMask", False)) + bpc = int(stream.get("/BitsPerComponent", 1 if mask else 8)) + n = 1 if mask else components(stream.get("/ColorSpace")) + mode = PNG_MODES.get((n, bpc)) + if mode is None: + return None + width, height = int(stream.Width), int(stream.Height) + if len(samples) != height * ((width * n * bpc + 7) // 8): + return None + img = Image.frombytes(mode, (width, height), samples) + if img.tobytes() != samples: + return None # Pillow zeroes the pad bits at the end of a 1-bit row; set ones must stay as they are + buf = io.BytesIO() + img.save(buf, "PNG", optimize=True) + png, pos, idat = buf.getvalue(), 8, [] + while pos < len(png): + size, kind = struct.unpack(">I4s", png[pos:pos + 8]) + if kind == b"IDAT": + idat.append(png[pos + 8:pos + 8 + size]) + pos += 12 + size + parms = pikepdf.Dictionary(Predictor=PNG_OPTIMUM, Colors=n, BitsPerComponent=bpc, Columns=width) + return b"".join(idat), parms + + +def rewrite_objects(pdf: pikepdf.Pdf) -> dict: + """Applies the rewrite to every stream of `pdf` in place; returns counts of what changed.""" + counts = {"decoded": 0, "images": 0, "jpegs": 0} + for obj in pdf.objects: + if not isinstance(obj, pikepdf.Stream): + continue + fs = filters(obj) + if not fs: + continue + if obj.get("/Subtype") == "/Image": + raw = obj.read_raw_bytes() + if fs == ["/FlateDecode"]: + warnings_before = len(pdf.get_warnings()) + samples = obj.read_bytes() + if len(pdf.get_warnings()) > warnings_before: + continue # qpdf could not decode it cleanly; keep the stream as the reader gets it today + got = png_flate(obj, samples) + if got is not None and len(got[0]) < len(raw): + obj.write(got[0], filter=pikepdf.Name.FlateDecode, decode_parms=got[1]) + counts["images"] += 1 + elif fs == ["/FlateDecode", "/DCTDecode"]: + p = obj.get("/DecodeParms") + if p is not None and not isinstance(p, pikepdf.Array): + continue # one dictionary for two filters is malformed; leave it to the reader + flate_parms, dct_parms = decode_parms(obj, 2) + if has_predictor(flate_parms): + continue + jpeg = inflate_whole(raw) + if jpeg is None: + continue # a bad checksum, a cut end or bytes after it: keep it as the reader gets it today + obj.write(jpeg, filter=pikepdf.Name.DCTDecode, decode_parms=dct_parms) + counts["jpegs"] += 1 + continue + if set(fs) <= GENERALIZED: + warnings_before = len(pdf.get_warnings()) + data = obj.read_bytes() + if len(pdf.get_warnings()) > warnings_before: + continue # qpdf could not decode it cleanly; keep the stream as the reader gets it today + obj.write(data) + counts["decoded"] += 1 + return counts + + +def save(pdf: pikepdf.Pdf) -> bytes: + buf = io.BytesIO() + pdf.save(buf, object_stream_mode=pikepdf.ObjectStreamMode.generate, compress_streams=False, + stream_decode_level=pikepdf.StreamDecodeLevel.none, deterministic_id=True, + fix_metadata_version=False) + return buf.getvalue() + + +def payload(stream: pikepdf.Object) -> tuple: + """What a reader gets from a stream: the fully decoded data when every filter is generalized; otherwise the data + its last (image) filter reads, with that filter's name and decode parameters.""" + fs = filters(stream) + if not fs: + return ("data", stream.read_raw_bytes()) + if set(fs) <= GENERALIZED: + try: + return ("data", stream.read_bytes()) + except pikepdf.PdfError: + return ("raw", tuple(fs), stream.read_raw_bytes()) # the rewrite leaves such a stream as it is + parms = decode_parms(stream, len(fs)) + head, last = fs[:-1], fs[-1] + if head == []: + data = stream.read_raw_bytes() + elif head == ["/FlateDecode"] and not has_predictor(parms[0]): + data = inflate(stream.read_raw_bytes()) + else: + return ("raw", tuple(fs), stream.read_raw_bytes()) + return (last, canonical(parms[-1]), data) + + +def canonical(obj, pages: dict = None, depth: int = 0): + """A plain-Python value for comparing small objects (destinations, actions, decode parameters) across two + files: a page reference becomes ("page", its index), any other indirect object its content.""" + if depth > 32: + return ("deep",) + if isinstance(obj, pikepdf.Object) and obj.is_indirect and pages is not None and obj.objgen in pages: + return ("page", pages[obj.objgen]) + if isinstance(obj, pikepdf.Stream): + # By what a reader gets, as compare_graphs compares one: the rewrite may change how it is stored. + return ("stream", {str(k): canonical(obj[k], pages, depth + 1) for k in sorted(obj.keys()) + if k not in REWRITABLE_KEYS}, payload(obj)) + if isinstance(obj, pikepdf.Dictionary): + return {str(k): canonical(obj[k], pages, depth + 1) for k in sorted(obj.keys())} + if isinstance(obj, pikepdf.Array): + return [canonical(x, pages, depth + 1) for x in obj] + if isinstance(obj, pikepdf.Object): + return ("bytes", bytes(obj)) if obj._type_code == pikepdf.ObjectType.string else str(obj) + return (type(obj).__name__, obj) + + +def named_destinations(pdf: pikepdf.Pdf, pages: dict) -> dict: + dests = {} + old_style = pdf.Root.get("/Dests") + if isinstance(old_style, pikepdf.Dictionary): + for k in old_style.keys(): + dests[("catalog", str(k))] = canonical(old_style[k], pages) + names = pdf.Root.get("/Names") + if isinstance(names, pikepdf.Dictionary) and isinstance(names.get("/Dests"), pikepdf.Dictionary): + for k, v in pikepdf.NameTree(names.Dests).items(): + dests[("tree", k)] = canonical(v, pages) + return dests + + +def outline(pdf: pikepdf.Pdf, pages: dict) -> list: + """[(depth, title, destination, action)] in reading order, walking /First and /Next.""" + out, seen = [], set() + root = pdf.Root.get("/Outlines") + stack = [(root.get("/First"), 0)] if isinstance(root, pikepdf.Dictionary) else [] + while stack: + item, depth = stack.pop() + if not isinstance(item, pikepdf.Dictionary) or item.objgen in seen: + continue + seen.add(item.objgen) + out.append((depth, canonical(item.get("/Title")), canonical(item.get("/Dest"), pages), + canonical(item.get("/A"), pages))) + stack.append((item.get("/Next"), depth)) + stack.append((item.get("/First"), depth + 1)) + return out + + +def page_index(pdf: pikepdf.Pdf) -> dict: + return {p.objgen: i for i, p in enumerate(pdf.pages)} + + +def same_scalar(a, b) -> bool: + if isinstance(a, pikepdf.Object) and isinstance(b, pikepdf.Object): + if a._type_code != b._type_code: + return False + if a._type_code == pikepdf.ObjectType.string: + return bytes(a) == bytes(b) + return a.unparse() == b.unparse() + return type(a) is type(b) and a == b + + +def compare_graphs(old: pikepdf.Pdf, new: pikepdf.Pdf) -> str | None: + """None when `new` is `old`'s object graph (from the trailer's /Root and /Info) with only the allowed changes; + otherwise where they first differ. A reference to a missing object reads as null on either side.""" + mapping, taken = {}, set() + stack = [(key, old.trailer.get(key), new.trailer.get(key)) for key in ("/Info", "/Root")] + while stack: + where, a, b = stack.pop() + a_ind = isinstance(a, pikepdf.Object) and a.is_indirect + b_ind = isinstance(b, pikepdf.Object) and b.is_indirect + if a_ind != b_ind: + return f"{where}: indirect on one side only" + if a_ind: + if a.objgen in mapping: + if mapping[a.objgen] != b.objgen: + return f"{where}: object {a.objgen} written as two objects" + continue + if b.objgen in taken: + return f"{where}: two objects written as one" + mapping[a.objgen] = b.objgen + taken.add(b.objgen) + if isinstance(a, pikepdf.Stream) or isinstance(b, pikepdf.Stream): + if not (isinstance(a, pikepdf.Stream) and isinstance(b, pikepdf.Stream)): + return f"{where}: a stream on one side only" + ka, kb = set(a.keys()) - REWRITABLE_KEYS, set(b.keys()) - REWRITABLE_KEYS + if ka != kb: + return f"{where}: stream keys differ ({sorted(ka ^ kb)})" + stack.extend((f"{where}{k}", a[k], b[k]) for k in sorted(ka)) + if payload(a) != payload(b): + return f"{where}: stream {a.objgen} decodes differently" + elif isinstance(a, pikepdf.Dictionary) or isinstance(b, pikepdf.Dictionary): + if not (isinstance(a, pikepdf.Dictionary) and isinstance(b, pikepdf.Dictionary)): + return f"{where}: a dictionary on one side only" + if set(a.keys()) != set(b.keys()): + return f"{where}: keys differ ({sorted(set(a.keys()) ^ set(b.keys()))})" + stack.extend((f"{where}{k}", a[k], b[k]) for k in sorted(a.keys())) + elif isinstance(a, pikepdf.Array) or isinstance(b, pikepdf.Array): + if not (isinstance(a, pikepdf.Array) and isinstance(b, pikepdf.Array)) or len(a) != len(b): + return f"{where}: arrays differ" + stack.extend((f"{where}[{i}]", x, y) for i, (x, y) in enumerate(zip(a, b))) + elif not same_scalar(a, b): + return f"{where}: {a!r} became {b!r}" + return None + + +def verify(old_data: bytes, new_data: bytes) -> None: + """Raises Refused unless `new_data` passes every check against `old_data`.""" + with pikepdf.open(io.BytesIO(old_data)) as old, pikepdf.open(io.BytesIO(new_data)) as new: + # The graph first: qpdf repairs a damaged page tree or name tree in memory when one is read (the checks + # below read both), and the comparison has to see each file as it is written. + difference = compare_graphs(old, new) + if difference: + raise Refused(difference, "graph") + problems = new.check_pdf_syntax() + if problems: + # qpdf names an in-memory file by its Python repr; say "the PDF" instead. + problem = re.sub(r"stream <_io\.BytesIO object at 0x[0-9a-f]+>", "the PDF", problems[0]) + raise Refused(f"the syntax check reports {problem}", "syntax") + old_pages, new_pages = page_index(old), page_index(new) + if len(old_pages) != len(new_pages): + raise Refused(f"{len(old_pages)} pages became {len(new_pages)}", "pages") + if named_destinations(old, old_pages) != named_destinations(new, new_pages): + raise Refused("the named destinations changed", "destinations") + if outline(old, old_pages) != outline(new, new_pages): + raise Refused("the outline changed", "outline") + + +def rewrite_pdf(data: bytes) -> tuple: + """(new PDF bytes, counts) for a PDF's bytes, checked; raises Refused when it cannot be rewritten safely, and + Skipped for a PDF this step leaves alone by policy.""" + try: + pdf = pikepdf.open(io.BytesIO(data)) + except pikepdf.PdfError as exc: + raise Refused(f"it cannot be opened ({exc})", "unreadable") from exc + with pdf: + if pdf.is_encrypted: + raise Skipped("it is encrypted", "encrypted") + # A signature covers the file's bytes, which any rewrite changes: the checks would pass and the signature + # would break. SigFlags bit 1 says the form holds a signature; /Perms holds DocMDP or usage-rights ones. + form = pdf.Root.get("/AcroForm") + if isinstance(form, pikepdf.Dictionary) and int(form.get("/SigFlags", 0)) & 1: + raise Skipped("it is signed (its /AcroForm sets SigFlags)", "signed") + if "/Perms" in pdf.Root: + raise Skipped("it is signed (it has a /Perms entry)", "signed") + counts = rewrite_objects(pdf) + new = save(pdf) + verify(data, new) + return new, counts + + +def keep_rewrite(old_size: int, new_size: int) -> bool: + """Whether a checked rewrite is worth writing, given both sizes as they ship.""" + return new_size < old_size + + +def rewrite_job(item: tuple) -> tuple: + """(path, new PDF bytes, counts), or (path, None, the LeftAsIs saying why), for the process pool.""" + path, data = item + try: + new, counts = rewrite_pdf(data) + except LeftAsIs as exc: + return path, None, exc + except Exception as exc: # noqa: BLE001 - anything else about one PDF leaves that PDF as it is + return path, None, Refused(f"it cannot be rewritten ({type(exc).__name__}: {exc})", "error") + return path, new, counts + + +def summary(rewritten: int, not_smaller: int, refused: Counter, skipped: Counter) -> str: + """One line: how many PDFs were rewritten, and why each of the others was left as it is.""" + def by_kind(counter): + return ", ".join(f"{kind} {n}" for kind, n in sorted(counter.items())) + parts = [f"{not_smaller} not smaller"] + if refused: + parts.append(f"{sum(refused.values())} refused by a check ({by_kind(refused)})") + if skipped: + parts.append(f"{sum(skipped.values())} skipped by policy ({by_kind(skipped)})") + return f"{rewritten} rewritten; left as they are: " + "; ".join(parts) + "." + + +def all_refused(refused: Counter, candidates: int) -> bool: + """Whether a check refused every PDF that policy did not skip: nothing is written, and the run fails.""" + return candidates > 0 and sum(refused.values()) == candidates + + +def describe(counts: dict) -> str: + return f"{counts['decoded']} stream(s) stored plain, {counts['images']} image(s) re-encoded, " \ + f"{counts['jpegs']} JPEG(s) unwrapped" + + +def run(cfg: dict) -> int: + db_path = cfg["db_path"] + if not db_path.is_file(): + print(f"error: {db_path} does not exist", file=sys.stderr) + return 1 + conn = sqlite3.connect(db_path) + codec = None + try: + lengths = {path: n for path, n in conn.execute("SELECT path, LENGTH(content) FROM Content")} + rows = [r for r in conn.execute( + "SELECT c.path, c.content, c.languageID, c.contentTypeID, c.templateId, ct.compression " + "FROM Content c JOIN ContentTypes ct ON ct.id = c.contentTypeID WHERE ct.value = ? ORDER BY c.path", + (PDF_TYPE,)) if not is_continuation_path(lengths, r[0])] + codec = BrotliCodec(load_dictionary(conn)) + work, problems, refused, skipped, not_smaller = [], [], Counter(), Counter(), 0 + for path, content, lang, ctid, tid, comp in rows: + if comp != "brotli": + print(f" {path}: skipped: stored with compression {comp!r}, not brotli") + skipped["not brotli"] += 1 + continue + blob = reassemble(conn, path, content) + try: + work.append((path, blob, codec.decompress(blob, comp), lang, ctid, tid)) + except Exception as exc: # noqa: BLE001 - one unreadable row stops the run + problems.append(f"{path}: cannot be read ({exc})") + if problems: + for p in problems: + print(f" error: {p}", file=sys.stderr) + print(f"error: {len(problems)} PDF row(s) cannot be read; nothing was written", file=sys.stderr) + return 1 + print(f"{len(rows)} PDF row(s); {len(work)} to try.") + + jobs = [(w[0], w[2]) for w in work] + if cfg["workers"] > 1: + with ProcessPoolExecutor(max_workers=cfg["workers"]) as pool: + done = list(pool.map(rewrite_job, jobs)) + else: + done = [rewrite_job(job) for job in jobs] + results = {path: (new, info) for path, new, info in done} + changes = [] + for path, blob, plain, lang, ctid, tid in work: + new, info = results[path] + if new is None: + print(f" {path}: left as it is: {info}") + (skipped if isinstance(info, Skipped) else refused)[info.kind] += 1 + continue + stored = codec.compress(new, "brotli") + if not keep_rewrite(len(blob), len(stored)): + not_smaller += 1 + if cfg["verbose"]: + print(f" {path}: left as it is: {human(len(stored))} stored is not smaller than " + f"{human(len(blob))}") + continue + print(f" {path}: {human(len(blob))} -> {human(len(stored))} stored; plain PDF {human(len(plain))} -> " + f"{human(len(new))}" + (f"; {describe(info)}" if cfg["verbose"] else "")) + changes.append((path, lengths[path], len(blob), stored, lang, ctid, tid)) + print(summary(len(changes), not_smaller, refused, skipped)) + candidates = len(work) - sum(1 for new, info in results.values() if isinstance(info, Skipped)) + if all_refused(refused, candidates): + print("error: a check refused every PDF it tried; nothing was written", file=sys.stderr) + return 1 + if not changes: + print("Nothing to do: no PDF gets smaller.") + return 0 + saved = sum(before - len(stored) for _p, _b, before, stored, *_ in changes) + print(f"{len(changes)} PDF row(s) to rewrite, saving {human(saved)} stored bytes.") + + if cfg["dry_run"]: + print("Dry run: nothing written.") + return 0 + print(f"Backup written to {backup_database(db_path)}") + conn.execute("BEGIN") + try: + # replace_row UPDATEs the base row, so the Bookshelf triggers (AFTER INSERT / AFTER DELETE on a '%.pdf' + # path) never fire for it; the "-N" continuation rows it trims or adds do not match '%.pdf'. + for path, base_len, _before, stored, lang, ctid, tid in changes: + replace_row(conn, path, base_len, stored, lang, ctid, tid, lambda m: print(m, file=sys.stderr)) + conn.commit() + except Exception: + conn.rollback() + raise + finally: + if codec is not None: + codec.close() + conn.close() + + vac = sqlite3.connect(db_path) + try: + vac.execute("VACUUM") + finally: + vac.close() + print(f"Done. Rewrote {len(changes)} PDF row(s).") + return 0 + + +def run_file(cfg: dict) -> int: + source, target = cfg["file"] + if not source.is_file(): + print(f"error: {source} does not exist", file=sys.stderr) + return 1 + if target.suffixes[-2:] != [".pdf", ".br"]: + print(f"error: {target} must end in .pdf.br", file=sys.stderr) + return 1 + raw = source.read_bytes() + if source.name.endswith(".br"): + try: + plain = brotli.decompress(raw) + except brotli.error as exc: + print(f"error: {source} is not plain Brotli ({exc})", file=sys.stderr) + return 1 + before = raw + else: + plain = raw + before = brotli.compress(plain, quality=FILE_QUALITY, lgwin=FILE_LGWIN) + _path, new, info = rewrite_job((str(source), plain)) + out, rewritten, refused, skipped, not_smaller = before, 0, Counter(), Counter(), 0 + if new is None: + print(f"{source}: left as it is: {info}") + (skipped if isinstance(info, Skipped) else refused)[info.kind] += 1 + else: + after = brotli.compress(new, quality=FILE_QUALITY, lgwin=FILE_LGWIN) + if keep_rewrite(len(before), len(after)): + out, rewritten = after, 1 + print(f"{source}: {human(len(before))} -> {human(len(after))} bytes; plain PDF {human(len(plain))} -> " + f"{human(len(new))}; {describe(info)}") + else: + not_smaller = 1 + print(f"{source}: left as it is: {human(len(after))} is not smaller than {human(len(before))}") + print(summary(rewritten, not_smaller, refused, skipped)) + if all_refused(refused, 1 - sum(skipped.values())): + print(f"error: a check refused {source}; nothing was written", file=sys.stderr) + return 1 + if cfg["dry_run"]: + print("Dry run: nothing written.") + return 0 + target.parent.mkdir(parents=True, exist_ok=True) + # The target keeps its mode; a new one gets what open() would give it. mkstemp's 0600 would otherwise stay. + if target.exists(): + mode = stat.S_IMODE(target.stat().st_mode) + else: + umask = os.umask(0) + os.umask(umask) + mode = 0o666 & ~umask + # Through a temporary file in the same directory, so the target is never half written and may be the source. + fd, tmp = tempfile.mkstemp(prefix=target.name + ".", dir=target.parent) + try: + with os.fdopen(fd, "wb") as fh: + fh.write(out) + os.chmod(tmp, mode) + os.replace(tmp, target) + except BaseException: + Path(tmp).unlink(missing_ok=True) + raise + print(f"Wrote {target}") + return 0 + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("db_path", type=Path, nargs="?", help="SQLite database to edit in place, e.g. documentation.db") + p.add_argument("--file", type=Path, nargs=2, metavar=("SOURCE", "TARGET"), + help="Rewrite one .pdf or plain-Brotli .pdf.br into TARGET (.pdf.br) instead of a database") + p.add_argument("--dry-run", action="store_true", help="Do all the work and checks; write nothing") + p.add_argument("--workers", type=int, default=8, help="Processes that rewrite PDFs; 1 runs in this one " + "(default: 8)") + p.add_argument("--verbose", action="store_true", help="Print what each rewrite changed") + return p + + +def main() -> None: + parser = build_parser() + args = parser.parse_args() + if (args.db_path is None) == (args.file is None): + parser.error("give a database or --file SOURCE TARGET, not both") + if args.file: + sys.exit(run_file({"file": args.file, "dry_run": args.dry_run})) + sys.exit(run({"db_path": args.db_path, "dry_run": args.dry_run, "workers": args.workers, + "verbose": args.verbose})) + + +if __name__ == "__main__": + main() diff --git a/scripts/test_optimize_pdfs.py b/scripts/test_optimize_pdfs.py new file mode 100644 index 00000000..b392a456 --- /dev/null +++ b/scripts/test_optimize_pdfs.py @@ -0,0 +1,1192 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.9" +# dependencies = [ +# "brotli", +# "Pillow", +# "pikepdf>=10", +# "pytest", +# ] +# /// +""" +test_optimize_pdfs.py + +Covers optimize_pdfs.py: what the rewrite changes in a PDF (and what it must +leave alone), the checks that refuse a bad rewrite, and whole runs over real +SQLite databases (Bookshelf triggers, chunked rows, dry run, rerun, rows that +are not Brotli, the database's own dictionary) and over files. + +Run it either way: + + uv run scripts/test_optimize_pdfs.py + uv run --with brotli --with Pillow --with pikepdf --with pytest pytest scripts/test_optimize_pdfs.py +""" +from __future__ import annotations + +import base64 +import io +import os +import random +import shutil +import sqlite3 +import stat +import sys +import tempfile +import time +import traceback +import zlib +from contextlib import redirect_stderr, redirect_stdout +from pathlib import Path + +import brotli +import pikepdf +import pikepdf.jbig2 +from pikepdf import Array, Dictionary, Name +from PIL import Image + +import optimize_pdfs as tool # noqa: E402 +from optimize_db_media import CHUNK_SIZE, BrotliCodec, reassemble # noqa: E402 +from update_media_references import replace_row # noqa: E402 + +CONTENT = (b"BT /F1 12 Tf 72 720 Td (" + b"Hello, offline world. " * 40 + b") Tj ET\n" + + b"q 64 0 0 64 72 72 cm /Rgb Do Q\n" * 8) + + +def gradient(mode: str, size=(64, 48)) -> Image.Image: + """Smooth samples: PNG predictors shrink these a lot against Flate with none.""" + img = Image.new("RGB", size) + img.putdata([(x * 4 % 256, y * 5 % 256, (x + y) * 2 % 256) for y in range(size[1]) for x in range(size[0])]) + return img.convert(mode) + + +def noisy_rows(width: int = 512, height: int = 64) -> Image.Image: + """1-bit rows that each repeat the one above with two bits flipped: the Up predictor turns them into zeros.""" + random.seed(9) + row, rows = bytearray(random.randrange(256) for _ in range(width // 8)), [] + for _ in range(height): + row = bytearray(row) + for _ in range(2): + row[random.randrange(len(row))] ^= 1 << random.randrange(8) + rows.append(bytes(row)) + return Image.frombytes("1", (width, height), b"".join(rows)) + + +def letters(n: int, seed: int) -> bytes: + """Random text over 16 letters: repeated, it is a few MB of PDF that Brotli at q11 shrinks to almost nothing.""" + random.seed(seed) + return bytes(random.choice(b"abcdefghijklmnop") for _ in range(n)) + + +def jpeg_bytes(size=(80, 60)) -> bytes: + random.seed(3) + img = Image.new("RGB", size) + img.putdata([(random.randrange(256), random.randrange(256), 128) for _ in range(size[0] * size[1])]) + buf = io.BytesIO() + img.save(buf, "JPEG", quality=70) + return buf.getvalue() + + +def image_stream(pdf: pikepdf.Pdf, img: Image.Image, colour_space, bpc: int, level: int = 6, + **extra) -> pikepdf.Stream: + """An image XObject holding `img`'s samples as Flate with no predictor, as many PDF producers write them.""" + s = pikepdf.Stream(pdf, zlib.compress(img.tobytes(), level)) + s.Type, s.Subtype, s.Width, s.Height = Name.XObject, Name.Image, img.width, img.height + s.BitsPerComponent = bpc + if colour_space is not None: + s.ColorSpace = colour_space + s.Filter = Name.FlateDecode + for k, v in extra.items(): + s["/" + k] = v + return pdf.make_indirect(s) + + +def sample_pdf(pages: int = 2, extra_content: bytes = b"") -> bytes: + """A small PDF with each kind of stream the rewrite treats differently, named destinations both ways, and an + outline with a destination and an action.""" + pdf = pikepdf.new() + for _ in range(pages): + pdf.add_blank_page(page_size=(612, 792)) + jpeg = jpeg_bytes() + images = { + "/Rgb": image_stream(pdf, gradient("RGB"), Name.DeviceRGB, 8), + "/Gray": image_stream(pdf, gradient("L"), Name.DeviceGray, 8), + "/Mask": image_stream(pdf, noisy_rows(), None, 1, level=1, ImageMask=True), # a fast-Flate producer + "/Bare": pdf.make_indirect(pikepdf.Stream(pdf, jpeg, Type=Name.XObject, Subtype=Name.Image, Width=80, + Height=60, BitsPerComponent=8, ColorSpace=Name.DeviceRGB, + Filter=Name.DCTDecode)), + "/Wrapped": pdf.make_indirect(pikepdf.Stream( + pdf, zlib.compress(jpeg, 9), Type=Name.XObject, Subtype=Name.Image, Width=80, Height=60, + BitsPerComponent=8, ColorSpace=Name.DeviceRGB, Filter=Array([Name.FlateDecode, Name.DCTDecode]), + DecodeParms=Array([None, Dictionary(ColorTransform=1)]))), + } + font_file = pdf.make_indirect(pikepdf.Stream(pdf, zlib.compress(b"\x00\x01fontdata" * 300), Filter=Name.FlateDecode, + Length1=2700)) + font = Dictionary(Type=Name.Font, Subtype=Name.TrueType, BaseFont=Name.Demo, + FontDescriptor=pdf.make_indirect(Dictionary(Type=Name.FontDescriptor, FontName=Name.Demo, + FontFile2=font_file))) + for page in pdf.pages: + page.obj.Contents = pdf.make_indirect(pikepdf.Stream(pdf, zlib.compress(CONTENT + extra_content, 9), + Filter=Name.FlateDecode)) + page.obj.Resources = Dictionary(XObject=Dictionary({k: v for k, v in images.items()}), + Font=Dictionary(F1=font)) + first = pdf.pages[0].obj + pdf.Root.Dests = Dictionary(old=Array([first, Name.Fit])) + tree = pikepdf.NameTree.new(pdf) + tree["chapter.1"] = Array([pdf.pages[-1].obj, Name.XYZ, 0, 700, None]) + tree["chapter.2"] = Array([first, Name.FitH, 500]) + pdf.Root.Names = Dictionary(Dests=tree.obj) + pdf.docinfo["/Title"] = pikepdf.String("A sample") + pdf.Root.Lang = pikepdf.String("en") + with pdf.open_outline() as o: + one = pikepdf.OutlineItem("One", 0) + one.children.append(pikepdf.OutlineItem("One.1", pages - 1)) + o.root.append(one) + o.root.append(pikepdf.OutlineItem("Two", action=Dictionary(S=Name.GoTo, D=Array([first, Name.Fit])))) + buf = io.BytesIO() + pdf.save(buf, compress_streams=False, object_stream_mode=pikepdf.ObjectStreamMode.disable, + stream_decode_level=pikepdf.StreamDecodeLevel.none, # keep each stream exactly as built + fix_metadata_version=False) + return buf.getvalue() + + +def _skip(reason: str) -> None: + """Skips the test: through pytest when it is running (so a skip is not reported as a pass), else by saying so.""" + if "pytest" in sys.modules: + import pytest + pytest.skip(reason) + print(f" (skipped: {reason})") + + +def open_pdf(data: bytes) -> pikepdf.Pdf: + return pikepdf.open(io.BytesIO(data)) + + +def rebuild(modify, base: bytes = None) -> bytes: + """sample_pdf() (or `base`) after `modify(pdf)`, every stream kept exactly as built.""" + pdf = open_pdf(base or sample_pdf()) + modify(pdf) + buf = io.BytesIO() + pdf.save(buf, compress_streams=False, stream_decode_level=pikepdf.StreamDecodeLevel.none, + fix_metadata_version=False) + return buf.getvalue() + + +def flate_stream(pdf: pikepdf.Pdf, data: bytes, **keys) -> pikepdf.Stream: + return pdf.make_indirect(pikepdf.Stream(pdf, zlib.compress(data, 9), Filter=Name.FlateDecode, **keys)) + + +def xobject(pdf: pikepdf.Pdf, name: str) -> pikepdf.Stream: + return pdf.pages[0].obj.Resources.XObject[name] + + +# --- the rewrite -------------------------------------------------------------------------------- + +def test_a_flate_content_stream_and_a_font_are_stored_plain(): + old = sample_pdf() + new, counts = tool.rewrite_pdf(old) + with open_pdf(old) as a, open_pdf(new) as b: + for i in range(2): + assert tool.filters(b.pages[i].obj.Contents) == [] + assert b.pages[i].obj.Contents.read_raw_bytes() == a.pages[i].obj.Contents.read_bytes() + font_file = b.pages[0].obj.Resources.Font.F1.FontDescriptor.FontFile2 + assert tool.filters(font_file) == [] and font_file.read_raw_bytes() == b"\x00\x01fontdata" * 300 + assert font_file.Length1 == 2700 # the font's own keys stay + assert counts["decoded"] >= 3 + + +def test_flate_images_get_png_predictors_and_give_back_the_same_samples(): + old = sample_pdf() + new, counts = tool.rewrite_pdf(old) + with open_pdf(old) as a, open_pdf(new) as b: + for name, colours, bpc in (("/Rgb", 3, 8), ("/Gray", 1, 8), ("/Mask", 1, 1)): + before, after = xobject(a, name), xobject(b, name) + assert tool.filters(after) == ["/FlateDecode"], name + assert after.DecodeParms.Predictor == tool.PNG_OPTIMUM, name + assert (after.DecodeParms.Colors, after.DecodeParms.BitsPerComponent, after.DecodeParms.Columns) == \ + (colours, bpc, int(before.Width)), name + assert after.read_bytes() == before.read_bytes(), name # every sample as it was + assert len(after.read_raw_bytes()) < len(before.read_raw_bytes()), name + assert counts["images"] == 3 + + +def test_a_flate_wrapped_jpeg_is_stored_bare_with_its_dct_parameters(): + old = sample_pdf() + new, counts = tool.rewrite_pdf(old) + with open_pdf(new) as b: + wrapped = xobject(b, "/Wrapped") + assert tool.filters(wrapped) == ["/DCTDecode"] + assert wrapped.read_raw_bytes() == jpeg_bytes() + assert wrapped.DecodeParms == Dictionary(ColorTransform=1) + assert counts["jpegs"] == 1 + + +def test_a_bare_jpeg_is_left_byte_for_byte(): + new, _counts = tool.rewrite_pdf(sample_pdf()) + with open_pdf(new) as b: + bare = xobject(b, "/Bare") + assert tool.filters(bare) == ["/DCTDecode"] and bare.read_raw_bytes() == jpeg_bytes() + assert "/DecodeParms" not in bare + + +def _with_image(build) -> tuple: + """(original bytes, rewritten bytes or None, counts or reason) for sample_pdf() plus one extra image `build` + makes, drawn on the first page as /Extra.""" + pdf = open_pdf(sample_pdf()) + pdf.pages[0].obj.Resources.XObject.Extra = pdf.make_indirect(build(pdf)) + buf = io.BytesIO() + pdf.save(buf, compress_streams=False, stream_decode_level=pikepdf.StreamDecodeLevel.none) + _path, new, info = tool.rewrite_job(("x.pdf", buf.getvalue())) + return buf.getvalue(), new, info + + +def test_a_wrapped_jpeg_whose_flate_has_a_predictor_is_left_as_it_is(): + jpeg = jpeg_bytes() + + def build(pdf): + rows = b"\x00" + jpeg # one row, PNG filter None: the predicted bytes are the JPEG plus a filter byte + return pikepdf.Stream(pdf, zlib.compress(rows), Type=Name.XObject, Subtype=Name.Image, Width=80, Height=60, + BitsPerComponent=8, ColorSpace=Name.DeviceRGB, + Filter=Array([Name.FlateDecode, Name.DCTDecode]), + DecodeParms=Array([Dictionary(Predictor=10, Columns=len(jpeg)), None])) + + old, new, counts = _with_image(build) + assert new is not None and counts["jpegs"] == 1 # the other wrapped JPEG still is unwrapped + with open_pdf(old) as a, open_pdf(new) as b: + before, after = xobject(a, "/Extra"), xobject(b, "/Extra") + assert tool.filters(after) == ["/FlateDecode", "/DCTDecode"] + assert after.read_raw_bytes() == before.read_raw_bytes() + + +def test_a_wrapped_jpeg_with_one_decode_dictionary_for_two_filters_is_left_as_it_is(): + # Malformed: readers differ on which filter the dictionary is for, so its meaning cannot be kept. + def build(pdf): + return pikepdf.Stream(pdf, zlib.compress(jpeg_bytes()), Type=Name.XObject, Subtype=Name.Image, Width=80, + Height=60, BitsPerComponent=8, ColorSpace=Name.DeviceRGB, + Filter=Array([Name.FlateDecode, Name.DCTDecode]), + DecodeParms=Dictionary(ColorTransform=0)) + + old, new, counts = _with_image(build) + assert new is not None and counts["jpegs"] == 1 + with open_pdf(new) as b: + after = xobject(b, "/Extra") + assert tool.filters(after) == ["/FlateDecode", "/DCTDecode"] and after.DecodeParms.ColorTransform == 0 + + +def test_an_encrypted_pdf_is_left_as_it_is(): + pdf = open_pdf(sample_pdf()) + buf = io.BytesIO() + pdf.save(buf, encryption=pikepdf.Encryption(owner="owner", user="", allow=pikepdf.Permissions(extract=False))) + _path, new, reason = tool.rewrite_job(("x.pdf", buf.getvalue())) + assert new is None and isinstance(reason, tool.Skipped) and str(reason) == "it is encrypted" + + +def test_a_signed_pdf_is_left_as_it_is(): + # A signature covers the file's bytes: every check would pass and the signature would break. + def modify(pdf): + pdf.Root.AcroForm = Dictionary(Fields=Array([]), SigFlags=3) + _path, new, reason = tool.rewrite_job(("x.pdf", rebuild(modify))) + assert new is None and isinstance(reason, tool.Skipped) and reason.kind == "signed" + assert str(reason) == "it is signed (its /AcroForm sets SigFlags)" + + +def test_a_pdf_with_signature_permissions_is_left_as_it_is(): + def modify(pdf): + pdf.Root.Perms = Dictionary(DocMDP=Dictionary(Type=Name.Sig)) + _path, new, reason = tool.rewrite_job(("x.pdf", rebuild(modify))) + assert new is None and isinstance(reason, tool.Skipped) and str(reason) == "it is signed (it has a /Perms entry)" + + +def test_a_form_without_a_signature_is_rewritten(): + def modify(pdf): + pdf.Root.AcroForm = Dictionary(Fields=Array([]), SigFlags=2) # append-only, no signature + _path, new, counts = tool.rewrite_job(("x.pdf", rebuild(modify))) + assert new is not None, counts + + +def test_an_image_the_predictor_does_not_shrink_is_left_as_it_is(): + pdf = pikepdf.new() + pdf.add_blank_page() + random.seed(11) + noise = Image.frombytes("L", (64, 64), bytes(random.randrange(256) for _ in range(64 * 64))) + s = image_stream(pdf, noise, Name.DeviceGray, 8) + pdf.pages[0].obj.Resources = Dictionary(XObject=Dictionary(N=s)) + raw = s.read_raw_bytes() + tool.rewrite_objects(pdf) + assert s.read_raw_bytes() == raw and "/DecodeParms" not in s + + +def test_named_destinations_and_the_outline_are_kept(): + old = sample_pdf() + new, _counts = tool.rewrite_pdf(old) + with open_pdf(new) as b: + pages = [p.obj.objgen for p in b.pages] + assert b.Root.Dests.old[0].objgen == pages[0] and b.Root.Dests.old[1] == Name.Fit + tree = pikepdf.NameTree(b.Root.Names.Dests) + assert sorted(tree.keys()) == ["chapter.1", "chapter.2"] + assert tree["chapter.1"][0].objgen == pages[1] and tree["chapter.2"][0].objgen == pages[0] + with b.open_outline() as o: + titles = [item.title for item in o.root] + nested = [item.title for item in o.root[0].children] + assert titles == ["One", "Two"] and nested == ["One.1"] + first = b.Root.Outlines.First + assert first.Dest[0].objgen == pages[0] and first.Next.A.D[0].objgen == pages[0] + assert first.First.Dest[0].objgen == pages[1] + + +def test_the_rewrite_is_deterministic_and_a_second_rewrite_changes_nothing(): + old = sample_pdf() + once, _ = tool.rewrite_pdf(old) + time.sleep(1.1) # qpdf's other /ID hashes the time to the second + again, _ = tool.rewrite_pdf(old) + assert once == again # the /ID comes from the content, not the clock + twice, counts = tool.rewrite_pdf(once) + assert twice == once and counts == {"decoded": 0, "images": 0, "jpegs": 0} + + +def test_a_pdf_with_a_stream_qpdf_cannot_decode_cleanly_is_left_as_it_is(): + # Its partial data, stored plain, would pass every check: the step must not quietly repair what a reader gets. + pdf = open_pdf(sample_pdf()) + pdf.pages[0].obj.Contents.write(zlib.compress(CONTENT, 9)[:-30], filter=Name.FlateDecode) + buf = io.BytesIO() + pdf.save(buf, compress_streams=False, stream_decode_level=pikepdf.StreamDecodeLevel.none) + path, new, reason = tool.rewrite_job(("x.pdf", buf.getvalue())) + assert new is None and reason.kind == "syntax" + assert str(reason).startswith("the syntax check reports WARNING: the PDF") + + +def _cut_checksum(data: bytes) -> bytes: + """Flate data without its adler32: complete, but qpdf warns that it ended early.""" + return zlib.compress(data, 6)[:-4] + + +def test_a_flate_image_qpdf_warns_on_is_not_rewritten(): + def modify(pdf): + image = pdf.pages[0].obj.Resources.XObject.Rgb + image.write(_cut_checksum(image.read_bytes()), filter=Name.FlateDecode) + _path, new, reason = tool.rewrite_job(("x.pdf", rebuild(modify))) + # Left encoded, the image fails the syntax check, so the whole PDF stays as it is. + assert new is None and reason.kind == "syntax", reason + + +def test_a_wrapped_jpeg_with_a_cut_checksum_is_not_unwrapped(): + def modify(pdf): + pdf.pages[0].obj.Resources.XObject.Wrapped.write(_cut_checksum(jpeg_bytes()), + filter=Array([Name.FlateDecode, Name.DCTDecode]), + decode_parms=Array([None, Dictionary(ColorTransform=1)])) + old = rebuild(modify) + _path, new, info = tool.rewrite_job(("x.pdf", old)) + # qpdf's syntax check does not mind a cut checksum before a JPEG filter, so the PDF is rewritten, this JPEG not. + assert new is not None and info["jpegs"] == 0, info + with open_pdf(old) as a, open_pdf(new) as b: + assert tool.filters(xobject(b, "/Wrapped")) == ["/FlateDecode", "/DCTDecode"] + assert xobject(b, "/Wrapped").read_raw_bytes() == xobject(a, "/Wrapped").read_raw_bytes() + + +def test_a_wrapped_jpeg_with_bytes_after_its_flate_end_is_not_unwrapped(): + def modify(pdf): + pdf.pages[0].obj.Resources.XObject.Wrapped.write(zlib.compress(jpeg_bytes(), 9) + b"trailing", + filter=Array([Name.FlateDecode, Name.DCTDecode]), + decode_parms=Array([None, Dictionary(ColorTransform=1)])) + old = rebuild(modify) + _path, new, info = tool.rewrite_job(("x.pdf", old)) + assert new is not None, info + with open_pdf(old) as a, open_pdf(new) as b: + assert tool.filters(xobject(b, "/Wrapped")) == ["/FlateDecode", "/DCTDecode"] + assert xobject(b, "/Wrapped").read_raw_bytes() == xobject(a, "/Wrapped").read_raw_bytes() + assert info["jpegs"] == 0 + + +def test_inflate_whole_wants_the_whole_stream(): + data = jpeg_bytes() + assert tool.inflate_whole(zlib.compress(data)) == data + assert tool.inflate_whole(zlib.compress(data)[:-4]) is None # no checksum + assert tool.inflate_whole(zlib.compress(data)[:-1] + b"\x00") is None # a wrong checksum + assert tool.inflate_whole(zlib.compress(data) + b"x") is None # bytes after the end + + +def test_a_one_bit_image_with_its_pad_bits_set_is_left_as_it_is(): + # 509 pixels make 64 bytes a row, 3 of its bits padding. Pillow would zero them; a reader may not. + random.seed(13) + row, rows = bytearray(random.randrange(256) for _ in range(64)), [] + for _ in range(64): + row = bytearray(row) + row[random.randrange(63)] ^= 1 << random.randrange(8) + row[63] |= 0b111 + rows.append(bytes(row)) + samples = b"".join(rows) + + def modify(pdf): + s = pdf.make_indirect(pikepdf.Stream(pdf, zlib.compress(samples, 1), Type=Name.XObject, Subtype=Name.Image, + Width=509, Height=64, ImageMask=True, BitsPerComponent=1, + Filter=Name.FlateDecode)) + pdf.pages[0].obj.Resources.XObject.Odd = s + old = rebuild(modify) + _path, new, info = tool.rewrite_job(("x.pdf", old)) + assert new is not None, info + with open_pdf(old) as a, open_pdf(new) as b: + assert xobject(b, "/Odd").read_raw_bytes() == xobject(a, "/Odd").read_raw_bytes() + + +def test_png_flate_declines_samples_of_the_wrong_length(): + pdf = pikepdf.new() + image = image_stream(pdf, gradient("RGB"), Name.DeviceRGB, 8) + samples = image.read_bytes() + assert tool.png_flate(image, samples) is not None + assert tool.png_flate(image, samples[:-1]) is None and tool.png_flate(image, samples + b"\x00") is None + + +def test_an_iccbased_rgb_image_gets_a_png_predictor(): + def modify(pdf): + icc = flate_stream(pdf, b"\x00" * 3000 + b"an ICC profile", N=3) + pdf.pages[0].obj.Resources.XObject.Icc = image_stream(pdf, gradient("RGB"), Array([Name.ICCBased, icc]), 8) + old = rebuild(modify) + new, _counts = tool.rewrite_pdf(old) + with open_pdf(old) as a, open_pdf(new) as b: + after = xobject(b, "/Icc") + assert after.DecodeParms.Predictor == tool.PNG_OPTIMUM and after.DecodeParms.Colors == 3 + assert after.read_bytes() == xobject(a, "/Icc").read_bytes() + + +def test_the_rewrite_writes_object_streams(): + new, _counts = tool.rewrite_pdf(sample_pdf()) + assert b"/Type /ObjStm" in new + + +def test_xmp_metadata_is_kept_byte_for_byte(): + # pikepdf's fix_metadata_version would rewrite pdf:PDFVersion to the new file version (object streams need + # 1.5), changing the stream. + xmp = (b'' + b'1.3' + b'') + + def modify(pdf): + pdf.Root.Metadata = pdf.make_indirect(pikepdf.Stream(pdf, xmp, Type=Name.Metadata, Subtype=Name.XML)) + new, _counts = tool.rewrite_pdf(rebuild(modify)) + with open_pdf(new) as b: + assert b.Root.Metadata.read_bytes() == xmp + + +def test_keep_rewrite_wants_strictly_smaller(): + assert tool.keep_rewrite(10, 9) and not tool.keep_rewrite(10, 10) and not tool.keep_rewrite(10, 11) + + +# --- the checks: each one refuses the fault it exists for ------------------------------------ + +def _with_fault(fault, *switch_off, source: bytes = None): + """Runs rewrite_pdf on `source` (sample_pdf() by default) with `fault(pdf)` applied after the rewrite, and the + named checks replaced by no-ops, so the check under test is the only one that can see the fault. Returns the + Refused message, or None.""" + real_objects = tool.rewrite_objects + saved = {name: getattr(tool, name) for name in switch_off} + + def faulty(pdf): + counts = real_objects(pdf) + fault(pdf) + return counts + + tool.rewrite_objects = faulty + for name in switch_off: + setattr(tool, name, {"compare_graphs": lambda old, new: None, + "named_destinations": lambda pdf, pages: {}, + "outline": lambda pdf, pages: []}[name]) + try: + tool.rewrite_pdf(source or sample_pdf()) + except tool.Refused as exc: + return str(exc) + finally: + tool.rewrite_objects = real_objects + for name, value in saved.items(): + setattr(tool, name, value) + return None + + +def test_the_graph_check_refuses_a_changed_dictionary_value(): + def fault(pdf): + pdf.pages[0].obj.Resources.Font.F1.BaseFont = Name.Other + refusal = _with_fault(fault) + assert refusal is not None and "/BaseFont" in refusal, refusal + + +def test_the_graph_check_refuses_a_changed_stream_value(): + def fault(pdf): + pdf.pages[0].obj.Resources.Font.F1.FontDescriptor.FontFile2.Length1 = 1 + refusal = _with_fault(fault) + assert refusal is not None and "/Length1" in refusal, refusal + + +def test_the_graph_check_refuses_an_added_stream_key(): + def fault(pdf): + pdf.pages[0].obj.Resources.Font.F1.FontDescriptor.FontFile2.Length2 = 5 + refusal = _with_fault(fault) + assert refusal is not None and refusal.endswith("stream keys differ (['/Length2'])"), refusal + + +def test_the_graph_check_refuses_an_added_dictionary_key(): + def fault(pdf): + pdf.pages[0].obj.Resources.Font.F1.Encoding = Name.WinAnsiEncoding + refusal = _with_fault(fault) + assert refusal is not None and refusal.endswith("/Font/F1: keys differ (['/Encoding'])"), refusal + + +def test_the_graph_check_refuses_an_appended_array_element(): + def fault(pdf): + pdf.Root.Dests.old.append(0) + refusal = _with_fault(fault) + assert refusal is not None and refusal.endswith("/Root/Dests/old: arrays differ"), refusal + + +def test_the_graph_check_walks_the_document_information(): + def fault(pdf): + pdf.docinfo["/Title"] = pikepdf.String("Another title") + refusal = _with_fault(fault) + assert refusal is not None and refusal.startswith("/Info/Title: "), refusal + + +def test_the_graph_check_compares_strings(): + def fault(pdf): + pdf.Root.Lang = pikepdf.String("fr") + refusal = _with_fault(fault) + assert refusal is not None and refusal.startswith("/Root/Lang: "), refusal + + +def test_the_graph_check_refuses_an_object_written_as_two(): + def fault(pdf): + font = pdf.pages[1].obj.Resources.Font.F1 + font.FontDescriptor = pdf.make_indirect(Dictionary({k: v for k, v in font.FontDescriptor.items()})) + refusal = _with_fault(fault) + assert refusal is not None and "written as two objects" in refusal, refusal + + +def test_the_graph_check_refuses_two_objects_written_as_one(): + def fault(pdf): + pdf.pages[1].obj.Contents = pdf.pages[0].obj.Contents # the same content, but one object instead of two + refusal = _with_fault(fault) + assert refusal is not None and refusal.endswith("two objects written as one"), refusal + + +def test_the_graph_check_refuses_a_stream_that_became_a_dictionary(): + def fault(pdf): + pdf.pages[0].obj.Contents = pdf.make_indirect(Dictionary(Length=0)) + refusal = _with_fault(fault) + assert refusal is not None and refusal.endswith("/Contents: a stream on one side only"), refusal + + +def test_the_graph_check_refuses_a_dictionary_that_became_an_array(): + def fault(pdf): + pdf.pages[0].obj.Resources.Font.F1 = Array([Name.Demo]) + refusal = _with_fault(fault) + assert refusal is not None and refusal.endswith("/Font/F1: a dictionary on one side only"), refusal + + +def test_the_graph_check_refuses_a_direct_object_made_indirect(): + def fault(pdf): + pdf.pages[0].obj.MediaBox = pdf.make_indirect(Array([0, 0, 612, 792])) + refusal = _with_fault(fault) + assert refusal is not None and refusal.endswith("/MediaBox: indirect on one side only"), refusal + + +def test_the_graph_check_refuses_changed_stream_data(): + def fault(pdf): + contents = pdf.pages[0].obj.Contents + contents.write(contents.read_bytes() + b" ") + refusal = _with_fault(fault) + assert refusal is not None and "decodes differently" in refusal, refusal + + +def test_the_graph_check_refuses_changed_image_samples(): + def fault(pdf): + image = pdf.pages[0].obj.Resources.XObject.Gray + samples = bytearray(image.read_bytes()) + samples[0] ^= 1 + image.write(zlib.compress(bytes(samples)), filter=Name.FlateDecode) + refusal = _with_fault(fault) + assert refusal is not None and "decodes differently" in refusal, refusal + + +def test_the_graph_check_refuses_a_changed_dct_parameter(): + def fault(pdf): + pdf.pages[0].obj.Resources.XObject.Wrapped.DecodeParms = Dictionary(ColorTransform=0) + refusal = _with_fault(fault) + assert refusal is not None and "decodes differently" in refusal, refusal + + +def test_the_page_count_check_refuses_a_lost_page(): + def fault(pdf): + del pdf.pages[-1] + assert _with_fault(fault, "compare_graphs", "named_destinations", "outline") == "2 pages became 1" + + +def test_the_named_destination_check_refuses_a_changed_destination(): + def fault(pdf): + pdf.Root.Dests.old = Array([pdf.pages[1].obj, Name.Fit]) + assert _with_fault(fault, "compare_graphs", "outline") == "the named destinations changed" + + +def test_the_named_destination_check_reads_the_name_tree(): + def fault(pdf): + pikepdf.NameTree(pdf.Root.Names.Dests)["chapter.2"] = Array([pdf.pages[1].obj, Name.Fit]) + assert _with_fault(fault, "compare_graphs", "outline") == "the named destinations changed" + + +def test_the_outline_check_refuses_a_changed_title(): + def fault(pdf): + pdf.Root.Outlines.First.Title = pikepdf.String("Changed") + assert _with_fault(fault, "compare_graphs", "named_destinations") == "the outline changed" + + +def test_the_outline_check_reads_nested_items(): + def fault(pdf): + pdf.Root.Outlines.First.First.Title = pikepdf.String("Changed") + assert _with_fault(fault, "compare_graphs", "named_destinations") == "the outline changed" + + +def test_the_outline_check_compares_actions(): + def fault(pdf): + pdf.Root.Outlines.First.Next.A.D = Array([pdf.pages[1].obj, Name.Fit]) + assert _with_fault(fault, "compare_graphs", "named_destinations") == "the outline changed" + + +def test_the_outline_check_compares_depth(): + def fault(pdf): + # One.1 moves from under One to between One and Two: the same items in the same order, one level up. + outlines = pdf.Root.Outlines + one = outlines.First + child, two = one.First, one.Next + for key in ("/First", "/Last", "/Count"): + if key in one: + del one[key] + one.Next, child.Prev, child.Next, child.Parent, two.Prev = child, one, two, outlines, child + assert _with_fault(fault, "compare_graphs", "named_destinations") == "the outline changed" + + +# A 64x16 JBIG2 image of four glyphs and its symbol dictionary, made by jbig2enc 0.31 (jbig2 -s -p): qpdf's +# syntax check decodes JBIG2, so the data has to be real. +JBIG2_GLOBALS = base64.b64decode("AAAAAAABAAAAACIAAAP//f8C/v7+AAAAAQAAAAFFD1Fw8qE5t2+8wnX/f/+s") +JBIG2_PAGE = base64.b64decode("AAAAATAAAQAAABMAAABAAAAAEAAAAAAAAAAAAAAAAAAAAgYiAAEAAAAdAAAAQAAAABAAAAAAAAAAAAAAAAAAAASf" + "AtfY/6w=") + + +def _with_jbig2_globals(pdf): + globals_ = flate_stream(pdf, JBIG2_GLOBALS) + pdf.pages[0].obj.Resources.XObject.Jb = pdf.make_indirect(pikepdf.Stream( + pdf, JBIG2_PAGE, Type=Name.XObject, Subtype=Name.Image, Width=64, Height=16, BitsPerComponent=1, + ColorSpace=Name.DeviceGray, Filter=Name.JBIG2Decode, DecodeParms=Dictionary(JBIG2Globals=globals_))) + + +def _with_script_action(pdf): + with pdf.open_outline() as o: + o.root.append(pikepdf.OutlineItem("Three", action=Dictionary( + S=Name.JavaScript, JS=flate_stream(pdf, b"app.alert('offline');" * 20)))) + + +def _with_stream_in_destination(pdf): + tree = pikepdf.NameTree(pdf.Root.Names.Dests) + tree["chapter.3"] = Dictionary(D=Array([pdf.pages[0].obj, Name.Fit]), Note=flate_stream(pdf, b"a note " * 50)) + + +def test_a_flate_jbig2_globals_stream_is_compared_by_its_data(): + try: + pikepdf.jbig2.get_decoder().check_available() + except Exception as exc: # noqa: BLE001 - DependencyError, or an OSError from looking for jbig2dec + return _skip(f"pikepdf cannot decode JBIG2 here ({exc}); the syntax check would refuse the fixture") + old = rebuild(_with_jbig2_globals) + new, _counts = tool.rewrite_pdf(old) # the globals stream is stored plain, and accepted + with open_pdf(new) as b: + assert tool.filters(xobject(b, "/Jb").DecodeParms.JBIG2Globals) == [] + + def fault(pdf): + g = pdf.pages[0].obj.Resources.XObject.Jb.DecodeParms.JBIG2Globals + g.write(g.read_bytes() + b"\x01") + refusal = _with_fault(fault, source=old) + assert refusal is not None and "decodes differently" in refusal, refusal + + +def test_an_outline_script_stream_is_compared_by_its_data(): + old = rebuild(_with_script_action) + tool.rewrite_pdf(old) + + def fault(pdf): + js = pdf.Root.Outlines.Last.A.JS + js.write(js.read_bytes() + b";") + assert _with_fault(fault, "compare_graphs", "named_destinations", source=old) == "the outline changed" + + +def test_a_stream_in_a_named_destination_is_compared_by_its_data(): + old = rebuild(_with_stream_in_destination) + tool.rewrite_pdf(old) + + def fault(pdf): + note = pikepdf.NameTree(pdf.Root.Names.Dests)["chapter.3"].Note + note.write(note.read_bytes() + b".") + assert _with_fault(fault, "compare_graphs", "outline", source=old) == "the named destinations changed" + + +def test_the_syntax_check_refuses_a_damaged_file(): + real_save = tool.save + tool.save = lambda pdf: real_save(pdf).replace(b"startxref", b"startxrex") + try: + tool.rewrite_pdf(sample_pdf()) + except tool.Refused as exc: + assert str(exc).startswith("the syntax check reports WARNING: the PDF: ") + else: + raise AssertionError("a damaged file was accepted") + finally: + tool.save = real_save + + +# --- whole runs over a database ------------------------------------------------------------------ + +PDF_ID, HTML_ID, RAW_PDF_ID = 14, 12, 30 +BOOK_PATH = "bookshelf/org.appdevforall.bookshelfplugin/Book.pdf" + + +def _database(tmp_path: Path, rows: dict, dictionary: bytes = None, raw_rows: dict = None, + fast: bool = False) -> Path: + """A database with the real Content/Bookshelf schema and its two triggers; every PDF row gets a curated + Bookshelf entry, as the shipped books have. `fast` stores the rows with Brotli's fastest settings (plain + databases only), so a few MB of PDF starts out chunked.""" + db_path = tmp_path / "documentation.db" + conn = sqlite3.connect(db_path) + conn.executescript( + "CREATE TABLE ContentTypes (id INTEGER PRIMARY KEY, value TEXT NOT NULL UNIQUE, compression TEXT NOT NULL);" + "CREATE TABLE Content (id INTEGER PRIMARY KEY AUTOINCREMENT, path TEXT NOT NULL, languageID INTEGER NOT NULL, " + "content BLOB NOT NULL, contentTypeID INTEGER NOT NULL, templateId INTEGER NOT NULL DEFAULT 0, UNIQUE(path));" + "CREATE TABLE Bookshelf (contentID INTEGER NOT NULL, title STRING DEFAULT '', description STRING DEFAULT '', " + "bookCategoryID INTEGER, UNIQUE(title, bookCategoryId));" + ) + conn.execute("INSERT INTO ContentTypes VALUES (?, 'application/pdf', 'brotli'), (?, 'text/html', 'brotli'), " + "(?, 'application/x-raw-pdf', 'none')", (PDF_ID, HTML_ID, RAW_PDF_ID)) + if dictionary is not None: + conn.execute("CREATE TABLE CompressionDictionary (id INTEGER PRIMARY KEY, data BLOB NOT NULL)") + conn.execute("INSERT INTO CompressionDictionary VALUES (1, ?)", (dictionary,)) + codec = BrotliCodec(dictionary) + for path, data in rows.items(): + type_id = HTML_ID if path.endswith(".html") else PDF_ID + conn.execute("INSERT INTO Content (path, languageID, content, contentTypeID, templateId) " + "VALUES (?, 1, x'00', ?, 7)", (path, type_id)) + stored = brotli.compress(data, quality=0, lgwin=10) if fast else codec.compress(data, "brotli") + replace_row(conn, path, 1, stored, 1, type_id, 7, lambda m: None) + for path, data in (raw_rows or {}).items(): + conn.execute("INSERT INTO Content (path, languageID, content, contentTypeID) VALUES (?, 1, ?, ?)", + (path, data, RAW_PDF_ID)) + codec.close() + conn.executescript( + "CREATE TRIGGER AddBook AFTER INSERT ON Content WHEN NEW.path LIKE '%.pdf' BEGIN " + "INSERT INTO Bookshelf (contentID, title) VALUES (NEW.id, CURRENT_TIMESTAMP || NEW.id); END;" + "CREATE TRIGGER DeleteBook AFTER DELETE ON Content WHEN OLD.path LIKE '%.pdf' BEGIN " + "DELETE FROM Bookshelf WHERE contentID = OLD.id; END;" + ) + for (content_id,) in conn.execute("SELECT id FROM Content WHERE path LIKE '%.pdf'").fetchall(): + conn.execute("INSERT INTO Bookshelf (contentID, title, description, bookCategoryID) VALUES (?, ?, ?, 2)", + (content_id, f"Curated {content_id}", "A curated description")) + conn.commit() + conn.close() + return db_path + + +def _read(db_path: Path, path: str, dictionary: bytes = None) -> bytes: + conn = sqlite3.connect(db_path) + codec = BrotliCodec(dictionary) + try: + content, comp = conn.execute("SELECT c.content, t.compression FROM Content c JOIN ContentTypes t " + "ON t.id = c.contentTypeID WHERE c.path = ?", (path,)).fetchone() + return codec.decompress(reassemble(conn, path, content), comp) + finally: + codec.close() + conn.close() + + +def _rows(db_path: Path) -> list: + conn = sqlite3.connect(db_path) + try: + return conn.execute("SELECT id, path, languageID, contentTypeID, templateId, LENGTH(content) FROM Content " + "ORDER BY path").fetchall() + finally: + conn.close() + + +def _bookshelf(db_path: Path) -> list: + conn = sqlite3.connect(db_path) + try: + return conn.execute("SELECT * FROM Bookshelf ORDER BY contentID").fetchall() + finally: + conn.close() + + +def _cfg(db_path: Path, **overrides) -> dict: + cfg = {"db_path": db_path, "dry_run": False, "workers": 1, "verbose": False} + cfg.update(overrides) + return cfg + + +def _two_rows() -> dict: + return {BOOK_PATH: sample_pdf(), "d/Guide.pdf": sample_pdf(pages=3), "d/page.html": b"guide"} + + +def test_end_to_end_rewrites_each_pdf_and_nothing_else(tmp_path: Path): + rows = _two_rows() + db_path = _database(tmp_path, rows) + before = {r[1]: r for r in _rows(db_path)} + assert tool.run(_cfg(db_path)) == 0 + after = {r[1]: r for r in _rows(db_path)} + for path in (BOOK_PATH, "d/Guide.pdf"): + new = _read(db_path, path) + assert new != rows[path] + tool.verify(rows[path], new) # the stored PDF passes every check against the original + assert after[path][:5] == before[path][:5] # same id, path, language, type, templateId + assert after[path][5] < before[path][5] + assert _read(db_path, "d/page.html") == rows["d/page.html"] + assert len(list(tmp_path.glob("*.backup-*"))) == 1 + conn = sqlite3.connect(db_path) + assert conn.execute("PRAGMA freelist_count").fetchone()[0] == 0 # VACUUMed after the rewrite + conn.close() + + +def test_the_bookshelf_rows_keep_their_ids_and_curated_titles(tmp_path: Path): + db_path = _database(tmp_path, _two_rows()) + shelf = _bookshelf(db_path) + assert len(shelf) == 2 and all(title.startswith("Curated") for _id, title, *_ in shelf) + assert tool.run(_cfg(db_path)) == 0 + assert _bookshelf(db_path) == shelf + + +def test_a_chunked_row_is_rejoined_and_its_continuations_trimmed(tmp_path: Path): + big = sample_pdf(extra_content=b"%" + letters(100_000, 5) * 30 + b"\n") + db_path = _database(tmp_path, {BOOK_PATH: big}, fast=True) + assert [r[1] for r in _rows(db_path)][:3] == [BOOK_PATH, f"{BOOK_PATH}-1", f"{BOOK_PATH}-2"] + shelf = _bookshelf(db_path) + assert tool.run(_cfg(db_path)) == 0 + assert [r[1] for r in _rows(db_path)] == [BOOK_PATH] + tool.verify(big, _read(db_path, BOOK_PATH)) + assert _bookshelf(db_path) == shelf + + +def test_a_row_still_over_a_chunk_is_written_back_in_chunks(tmp_path: Path): + random.seed(6) + buf = io.BytesIO() + # A 1.2 MB JPEG of noise: the rewrite leaves it as it is and no compressor shrinks it. + Image.frombytes("RGB", (720, 720), random.randbytes(720 * 720 * 3)).save(buf, "JPEG", quality=95, subsampling=0) + pdf = open_pdf(sample_pdf(extra_content=b"%" + letters(100_000, 5) * 30 + b"\n")) + pdf.Root.Noise = pdf.make_indirect(pikepdf.Stream(pdf, buf.getvalue(), Filter=Name.DCTDecode, Type=Name.XObject, + Subtype=Name.Image, Width=720, Height=720, + BitsPerComponent=8, ColorSpace=Name.DeviceRGB)) + buf = io.BytesIO() + pdf.save(buf, compress_streams=False, stream_decode_level=pikepdf.StreamDecodeLevel.none) + big = buf.getvalue() + db_path = _database(tmp_path, {BOOK_PATH: big}, fast=True) + assert len(_rows(db_path)) > 3 + assert tool.run(_cfg(db_path)) == 0 + rows = _rows(db_path) + assert [r[1] for r in rows] == [BOOK_PATH, f"{BOOK_PATH}-1"] + assert rows[0][5] == CHUNK_SIZE and rows[1][4] == 7 # the continuation keeps the base row's templateId + tool.verify(big, _read(db_path, BOOK_PATH)) + + +def test_dry_run_writes_nothing(tmp_path: Path): + db_path = _database(tmp_path, _two_rows()) + before = db_path.read_bytes() + out = io.StringIO() + with redirect_stdout(out): + assert tool.run(_cfg(db_path, dry_run=True)) == 0 + assert "Dry run: nothing written." in out.getvalue() + assert db_path.read_bytes() == before and not list(tmp_path.glob("*.backup-*")) + + +def test_a_second_run_changes_nothing(tmp_path: Path): + db_path = _database(tmp_path, _two_rows()) + assert tool.run(_cfg(db_path)) == 0 + once = db_path.read_bytes() + out = io.StringIO() + with redirect_stdout(out): + assert tool.run(_cfg(db_path)) == 0 + assert "Nothing to do" in out.getvalue() + assert db_path.read_bytes() == once and len(list(tmp_path.glob("*.backup-*"))) == 1 + + +def test_a_rewrite_that_is_not_smaller_is_left_alone(tmp_path: Path): + rows = _two_rows() + db_path = _database(tmp_path, rows) + before = db_path.read_bytes() + real = tool.rewrite_pdf + + def padded(data): + new, counts = real(data) + # Valid, checked, but larger once compressed: random bytes after %%EOF are ignored by readers. + return new + bytes(random.randrange(256) for _ in range(len(data))), counts + + tool.rewrite_pdf = padded + try: + assert tool.run(_cfg(db_path)) == 0 + finally: + tool.rewrite_pdf = real + assert db_path.read_bytes() == before + + +def _run_with_fault(db_path: Path, when) -> tuple: + """(exit code, stdout, stderr) of a run whose rewrite changes /BaseFont in the PDFs `when(pdf)` picks.""" + real = tool.rewrite_objects + + def faulty(pdf): + counts = real(pdf) + if when(pdf): + pdf.pages[0].obj.Resources.Font.F1.BaseFont = Name.Other + return counts + + tool.rewrite_objects = faulty + out, err = io.StringIO(), io.StringIO() + try: + with redirect_stdout(out), redirect_stderr(err): + code = tool.run(_cfg(db_path)) + finally: + tool.rewrite_objects = real + return code, out.getvalue(), err.getvalue() + + +def test_a_run_where_a_check_refuses_every_pdf_fails_with_nothing_written(tmp_path: Path): + rows = _two_rows() + db_path = _database(tmp_path, rows) + before = db_path.read_bytes() + code, out, err = _run_with_fault(db_path, lambda pdf: True) + assert code == 1 + assert f"{BOOK_PATH}: left as it is: /Root/Pages/Kids[0]/Resources/Font/F1/BaseFont" in out + assert "0 rewritten; left as they are: 0 not smaller; 2 refused by a check (graph 2)." in out + assert "a check refused every PDF it tried; nothing was written" in err + assert db_path.read_bytes() == before and not list(tmp_path.glob("*.backup-*")) + + +def test_a_run_where_a_check_refuses_some_pdfs_reports_them(tmp_path: Path): + rows = _two_rows() + db_path = _database(tmp_path, rows) + code, out, _err = _run_with_fault(db_path, lambda pdf: len(pdf.pages) == 3) # only d/Guide.pdf + assert code == 0 + assert "1 rewritten; left as they are: 0 not smaller; 1 refused by a check (graph 1)." in out + assert _read(db_path, "d/Guide.pdf") == rows["d/Guide.pdf"] and _read(db_path, BOOK_PATH) != rows[BOOK_PATH] + + +def test_a_run_fails_when_a_check_refuses_every_pdf_policy_did_not_skip(tmp_path: Path): + signed = rebuild(lambda pdf: setattr(pdf.Root, "AcroForm", Dictionary(Fields=Array([]), SigFlags=1))) + db_path = _database(tmp_path, {BOOK_PATH: signed, "d/Guide.pdf": sample_pdf(pages=3)}) + before = db_path.read_bytes() + code, out, _err = _run_with_fault(db_path, lambda pdf: True) + assert code == 1 + assert "0 rewritten; left as they are: 0 not smaller; 1 refused by a check (graph 1); " \ + "1 skipped by policy (signed 1)." in out + assert db_path.read_bytes() == before + + +def test_a_run_where_policy_skips_every_pdf_is_not_a_failure(tmp_path: Path): + signed = rebuild(lambda pdf: setattr(pdf.Root, "AcroForm", Dictionary(Fields=Array([]), SigFlags=1))) + db_path = _database(tmp_path, {BOOK_PATH: signed}) + before = db_path.read_bytes() + out = io.StringIO() + with redirect_stdout(out): + assert tool.run(_cfg(db_path)) == 0 + assert "0 rewritten; left as they are: 0 not smaller; 1 skipped by policy (signed 1)." in out.getvalue() + assert db_path.read_bytes() == before + + +def test_a_pdf_row_that_is_not_brotli_is_skipped_and_named(tmp_path: Path): + db_path = _database(tmp_path, {}, raw_rows={"x/raw.pdf": sample_pdf()}) + before = db_path.read_bytes() + out = io.StringIO() + with redirect_stdout(out): + real_type = tool.PDF_TYPE + tool.PDF_TYPE = "application/x-raw-pdf" + try: + assert tool.run(_cfg(db_path)) == 0 + finally: + tool.PDF_TYPE = real_type + assert "x/raw.pdf: skipped: stored with compression 'none', not brotli" in out.getvalue() + assert "0 rewritten; left as they are: 0 not smaller; 1 skipped by policy (not brotli 1)." in out.getvalue() + assert db_path.read_bytes() == before + + +def test_an_unreadable_row_stops_the_run_with_nothing_written(tmp_path: Path): + db_path = _database(tmp_path, _two_rows()) + conn = sqlite3.connect(db_path) + conn.execute("UPDATE Content SET content = x'0102030405' WHERE path = 'd/Guide.pdf'") + conn.commit() + conn.close() + before = db_path.read_bytes() + assert tool.run(_cfg(db_path)) == 1 + assert db_path.read_bytes() == before + + +def test_the_process_pool_gives_the_same_result(tmp_path: Path): + rows = _two_rows() # one set of rows for both: each sample_pdf() carries its own random /ID + (tmp_path / "a").mkdir() + (tmp_path / "b").mkdir() + one, two = _database(tmp_path / "a", rows), _database(tmp_path / "b", rows) + assert tool.run(_cfg(one, workers=1)) == 0 + assert tool.run(_cfg(two, workers=3)) == 0 + for path in (BOOK_PATH, "d/Guide.pdf"): + assert _read(one, path) == _read(two, path) + + +def test_rows_stay_in_the_database_dictionary(tmp_path: Path): + if shutil.which("brotli") is None: + return _skip("no brotli CLI") + dictionary = b"/Type /XObject /Subtype /Image /Filter /FlateDecode /DecodeParms BT ET Tj Tf " * 60 + rows = {BOOK_PATH: sample_pdf()} + db_path = _database(tmp_path, rows, dictionary=dictionary) + assert tool.run(_cfg(db_path)) == 0 + new = _read(db_path, BOOK_PATH, dictionary) + tool.verify(rows[BOOK_PATH], new) + conn = sqlite3.connect(db_path) + stored = conn.execute("SELECT content FROM Content WHERE path = ?", (BOOK_PATH,)).fetchone()[0] + conn.close() + try: + plain = brotli.decompress(stored) + except brotli.error: + plain = None + assert plain != new # it needs the dictionary to decode back + + +# --- file mode ----------------------------------------------------------------------------------- + +def _file_cfg(source: Path, target: Path, **overrides) -> dict: + cfg = {"file": (source, target), "dry_run": False} + cfg.update(overrides) + return cfg + + +def test_file_mode_takes_a_pdf(tmp_path: Path): + source = tmp_path / "Book.pdf" + source.write_bytes(sample_pdf()) + target = tmp_path / "out" / "Book.pdf.br" + assert tool.run_file(_file_cfg(source, target)) == 0 + new = brotli.decompress(target.read_bytes()) + tool.verify(source.read_bytes(), new) + as_shipped = brotli.compress(source.read_bytes(), quality=11, lgwin=24) + assert len(target.read_bytes()) < len(as_shipped) + assert target.read_bytes() == brotli.compress(new, quality=11, lgwin=24) + + +def test_file_mode_takes_a_plain_brotli_pdf_and_may_write_over_it(tmp_path: Path): + source = tmp_path / "Book.pdf.br" + source.write_bytes(brotli.compress(sample_pdf(), quality=11, lgwin=24)) + old = source.read_bytes() + assert tool.run_file(_file_cfg(source, source)) == 0 + new = brotli.decompress(source.read_bytes()) + tool.verify(brotli.decompress(old), new) + assert len(source.read_bytes()) < len(old) + assert [p.name for p in tmp_path.iterdir()] == ["Book.pdf.br"] # no temporary file left behind + assert tool.run_file(_file_cfg(source, source)) == 0 + assert brotli.decompress(source.read_bytes()) == new # a second run changes nothing + + +def _run_file_with(rewrite, source: Path, target: Path) -> tuple: + """(exit code, stdout, stderr) of a file-mode run with rewrite_pdf replaced by `rewrite`.""" + real = tool.rewrite_pdf + tool.rewrite_pdf = rewrite + out, err = io.StringIO(), io.StringIO() + try: + with redirect_stdout(out), redirect_stderr(err): + code = tool.run_file(_file_cfg(source, target)) + finally: + tool.rewrite_pdf = real + return code, out.getvalue(), err.getvalue() + + +def _book(tmp_path: Path) -> Path: + source = tmp_path / "Book.pdf.br" + source.write_bytes(brotli.compress(sample_pdf(), quality=11, lgwin=24)) + return source + + +def test_file_mode_writes_nothing_when_a_check_refuses(tmp_path: Path): + source, target = _book(tmp_path), tmp_path / "Out.pdf.br" + + def refuse(data): + raise tool.Refused("a test refusal", "graph") + + code, out, err = _run_file_with(refuse, source, target) + assert code == 1 and not target.exists() + assert "left as it is: a test refusal" in out + assert "0 rewritten; left as they are: 0 not smaller; 1 refused by a check (graph 1)." in out + assert "a check refused" in err + + +def test_file_mode_copies_the_input_when_policy_skips_it(tmp_path: Path): + source, target = _book(tmp_path), tmp_path / "Out.pdf.br" + + def skip(data): + raise tool.Skipped("it is signed (it has a /Perms entry)", "signed") + + code, out, _err = _run_file_with(skip, source, target) + assert code == 0 and target.read_bytes() == source.read_bytes() + assert "1 skipped by policy (signed 1)" in out + + +def test_file_mode_keeps_the_input_when_the_rewrite_is_not_smaller(tmp_path: Path): + source, target = _book(tmp_path), tmp_path / "Out.pdf.br" + real = tool.rewrite_pdf + + def padded(data): + new, counts = real(data) + return new + random.randbytes(len(data)), counts # ignored after %%EOF, but larger once compressed + + code, out, _err = _run_file_with(padded, source, target) + assert code == 0 and target.read_bytes() == source.read_bytes() + assert "0 rewritten; left as they are: 1 not smaller." in out + + +def test_file_mode_gives_a_new_file_the_mode_open_would(tmp_path: Path): + source, target = _book(tmp_path), tmp_path / "Out.pdf.br" + old_umask = os.umask(0o027) + try: + assert tool.run_file(_file_cfg(source, target)) == 0 + finally: + os.umask(old_umask) + assert stat.S_IMODE(target.stat().st_mode) == 0o640 + + +def test_file_mode_keeps_the_mode_of_the_file_it_replaces(tmp_path: Path): + source = _book(tmp_path) + source.chmod(0o604) + assert tool.run_file(_file_cfg(source, source)) == 0 + assert stat.S_IMODE(source.stat().st_mode) == 0o604 + + +def test_file_mode_dry_run_writes_nothing(tmp_path: Path): + source = tmp_path / "Book.pdf" + source.write_bytes(sample_pdf()) + target = tmp_path / "Book.pdf.br" + assert tool.run_file(_file_cfg(source, target, dry_run=True)) == 0 + assert not target.exists() + + +def test_file_mode_refuses_a_target_that_is_not_pdf_br(tmp_path: Path): + source = tmp_path / "Book.pdf" + source.write_bytes(sample_pdf()) + assert tool.run_file(_file_cfg(source, tmp_path / "Book.br")) == 1 + assert not (tmp_path / "Book.br").exists() + + +def _main() -> int: + failures = [] + for name, function in sorted(globals().items()): + if not name.startswith("test_") or not callable(function): + continue + with tempfile.TemporaryDirectory() as directory: + needs_tmp = "tmp_path" in function.__code__.co_varnames[:function.__code__.co_argcount] + try: + function(Path(directory)) if needs_tmp else function() + except (KeyboardInterrupt, SystemExit): + raise + except BaseException as exc: # noqa: BLE001 - a failure is the result here + frame = traceback.extract_tb(exc.__traceback__)[-1] + failures.append(f"{name} (line {frame.lineno}: {frame.line}): {type(exc).__name__}: {exc}") + print(f"FAIL {name}", file=sys.stderr) + else: + print(f"ok {name}") + print() + print(f"{len(failures)} failure(s)" if failures else "all tests passed") + for failure in failures: + print(f" {failure}", file=sys.stderr) + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(_main())