CoolFace
Apppublic

internationalscholarsprogram/handbook-engine

sourceHugging Faceupdated 5mo agoView on Hugging Face
0likes
pdf_renderer.py327 linesDownload Raw Back to services
1"""Playwright-based PDF renderer — Chromium headless PDF export.2 3Replaces WeasyPrint. Uses Playwright to launch headless Chromium,4load the fully-rendered HTML, wait for fonts/images/layout, and5export a print-quality PDF.6"""7 8from __future__ import annotations9 10import asyncio11import logging12import os13import tempfile14from pathlib import Path15from typing import Optional16 17logger = logging.getLogger(__name__)18 19# Singleton browser instance for reuse across requests20_browser = None21_browser_lock = asyncio.Lock()22 23 24async def _get_browser():25    """Get or create a persistent Chromium browser instance.26 27    Uses ``channel="chrome"`` so Playwright drives the system-installed28    Google Chrome (or Chromium) instead of requiring a separate browser29    download from the Playwright CDN.  Falls back to the default30    bundled Chromium if the system browser is not found.31    """32    global _browser33    async with _browser_lock:34        if _browser is None or not _browser.is_connected():35            from playwright.async_api import async_playwright36 37            pw = await async_playwright().start()38 39            launch_args = [40                "--no-sandbox",41                "--disable-setuid-sandbox",42                "--disable-dev-shm-usage",43                "--disable-gpu",44                "--font-render-hinting=none",45            ]46 47            # Try system Chrome first, then fall back to bundled Chromium48            try:49                _browser = await pw.chromium.launch(50                    channel="chrome",51                    headless=True,52                    args=launch_args,53                )54                logger.info("System Chrome launched for PDF rendering")55            except Exception:56                logger.warning(57                    "System Chrome not available, falling back to bundled Chromium"58                )59                _browser = await pw.chromium.launch(60                    headless=True,61                    args=launch_args,62                )63                logger.info("Bundled Chromium launched for PDF rendering")64        return _browser65 66 67async def shutdown_browser():68    """Gracefully close the browser on application shutdown."""69    global _browser70    async with _browser_lock:71        if _browser and _browser.is_connected():72            await _browser.close()73            _browser = None74            logger.info("Chromium browser closed")75 76 77async def render_pdf_from_html(78    html_content: str,79    *,80    format: str = "A4",81    print_background: bool = True,82    prefer_css_page_size: bool = True,83    wait_timeout: int = 30000,84) -> bytes:85    """Render HTML string to PDF bytes using Playwright Chromium.86 87    Generates a base PDF (content only, no decorative header/label),88    then creates a one-page overlay with the header image and right-side89    label, and stamps the overlay onto content pages (page 3 → last90    content page) using pypdf.  Pages 1-2 (cover/TOC) and trailing91    full-page image pages get no overlay.92 93    Args:94        html_content: Complete HTML document string.95        format: Page format (default A4).96        print_background: Include background colors/images.97        prefer_css_page_size: Use @page CSS rules for sizing.98        wait_timeout: Max time (ms) to wait for page load.99 100    Returns:101        PDF file bytes.102    """103    browser = await _get_browser()104    context = await browser.new_context(105        viewport={"width": 794, "height": 1123},  # A4 at 96dpi106        device_scale_factor=1,107        java_script_enabled=True,108    )109    page = await context.new_page()110 111    try:112        # Write HTML to a temp file so Chromium can load local file:// resources113        with tempfile.NamedTemporaryFile(114            mode="w",115            suffix=".html",116            delete=False,117            encoding="utf-8",118        ) as tmp:119            tmp.write(html_content)120            tmp_path = tmp.name121 122        try:123            file_url = Path(tmp_path).as_uri()124            await page.goto(file_url, wait_until="load", timeout=wait_timeout)125 126            # Wait for fonts and images to be fully loaded127            await page.evaluate("() => document.fonts.ready")128            await page.evaluate("""129                () => {130                    const images = Array.from(document.querySelectorAll('img'));131                    return Promise.all(images.map(img => {132                        if (img.complete) return Promise.resolve();133                        return new Promise(r => {134                            img.addEventListener('load', r);135                            img.addEventListener('error', r);136                        });137                    }));138                }139            """)140 141            # ── Collect info from DOM before hiding elements ──142            header_src = await page.evaluate("""143                () => {144                    const img = document.querySelector('.page-header img');145                    return img ? img.src : '';146                }147            """)148            label_src = await page.evaluate("""149                () => {150                    const img = document.querySelector('.hb-right-label img');151                    return img ? img.src : '';152                }153            """)154            num_bottom_pages = await page.evaluate("""155                () => document.querySelectorAll('.fullpage-img-wrap').length156            """)157            # Cover page count: cover + TOC image (each is a .cover-page)158            num_cover_pages = await page.evaluate("""159                () => document.querySelectorAll('.cover-page').length160            """)161 162            logger.info(163                "Overlay info: header=%s, label=%s, covers=%d, bottoms=%d",164                bool(header_src), bool(label_src),165                num_cover_pages, num_bottom_pages,166            )167 168            # ── Hide header, footer, and label from the base PDF ──169            await page.evaluate("""170                () => {171                    document.querySelectorAll('.page-header, .page-footer, .hb-right-label')172                        .forEach(el => el.style.display = 'none');173                }174            """)175 176            # ── Render BASE PDF (no header, no label) ──177            base_pdf = await page.pdf(178                format=format,179                print_background=print_background,180                prefer_css_page_size=prefer_css_page_size,181                margin={182                    "top": "2.54cm",183                    "right": "2.54cm",184                    "bottom": "2.54cm",185                    "left": "2.54cm",186                },187                display_header_footer=True,188                header_template='<span></span>',189                footer_template=(190                    '<div style="width:100%;text-align:center;font-size:9px;'191                    'font-family:Century Gothic,Segoe UI,sans-serif;color:#0263A3;'192                    'padding:0 0 6px 0;letter-spacing:0.5px;">'193                    '<span style="font-weight:700;" class="pageNumber"></span></div>'194                ),195            )196            logger.info("Base PDF rendered, size=%d bytes", len(base_pdf))197 198        finally:199            os.unlink(tmp_path)200 201        # ── Build overlay (header + label) and stamp onto content pages ──202        if not header_src and not label_src:203            logger.info("No header or label to overlay, returning base PDF")204            return base_pdf205 206        overlay_pdf = await _build_overlay_pdf(207            page, header_src, label_src, format, wait_timeout208        )209 210        merged = _stamp_overlay(211            base_pdf, overlay_pdf,212            skip_front=num_cover_pages,213            skip_back=num_bottom_pages,214        )215        logger.info("Final PDF with overlay, size=%d bytes", len(merged))216        return merged217 218    finally:219        await context.close()220 221 222async def _build_overlay_pdf(223    page, header_src: str, label_src: str,224    format: str, timeout: int,225) -> bytes:226    """Render a single-page transparent overlay PDF with header + label."""227    parts = []228    if header_src:229        parts.append(230            f'<div style="position:fixed;top:0;left:0;width:100%;height:2.54cm;'231            f'margin:0;padding:0;overflow:hidden;z-index:1;">'232            f'<img src="{header_src}" style="display:block;width:100%;'233            f'height:100%;object-fit:fill;margin:0;padding:0;" /></div>'234        )235    if label_src:236        # Word doc: 3.0cm × 22.7cm container, ~0.35cm bleeds past right edge.237        # Scaled proportionally to A4: 2.9cm × 24.1cm, right:-0.3cm to238        # let part bleed off-page just like the Word original.239        parts.append(240            f'<div style="position:fixed;top:3.5cm;right:-0.3cm;width:2.9cm;'241            f'height:24.1cm;z-index:2;overflow:visible;">'242            f'<img src="{label_src}" style="display:block;width:100%;'243            f'height:100%;object-fit:fill;" /></div>'244        )245 246    overlay_html = (247        '<!doctype html><html><head><meta charset="utf-8">'248        '<style>'249        '@page{size:A4;margin:0}'250        'html,body{margin:0;padding:0;background:transparent}'251        '</style></head><body>'252        + '\n'.join(parts)253        + '<div style="height:297mm;width:210mm;"></div>'254        '</body></html>'255    )256 257    with tempfile.NamedTemporaryFile(258        mode="w", suffix=".html", delete=False, encoding="utf-8",259    ) as tmp:260        tmp.write(overlay_html)261        tmp_path = tmp.name262 263    try:264        await page.goto(265            Path(tmp_path).as_uri(),266            wait_until="load",267            timeout=timeout,268        )269        await page.evaluate("() => document.fonts.ready")270        await page.evaluate("""271            () => {272                const images = Array.from(document.querySelectorAll('img'));273                return Promise.all(images.map(img => {274                    if (img.complete) return Promise.resolve();275                    return new Promise(r => {276                        img.addEventListener('load', r);277                        img.addEventListener('error', r);278                    });279                }));280            }281        """)282 283        overlay_bytes = await page.pdf(284            format=format,285            print_background=True,286            prefer_css_page_size=True,287            margin={"top": "0", "right": "0", "bottom": "0", "left": "0"},288            display_header_footer=False,289        )290        logger.info("Overlay PDF rendered, size=%d bytes", len(overlay_bytes))291        return overlay_bytes292    finally:293        os.unlink(tmp_path)294 295 296def _stamp_overlay(297    base_pdf: bytes,298    overlay_pdf: bytes,299    skip_front: int = 2,300    skip_back: int = 4,301) -> bytes:302    """Merge overlay onto content pages of the base PDF.303 304    Pages 0..(skip_front-1) and (total-skip_back)..(total-1) are left305    untouched.  All other pages get the overlay stamped on top.306    """307    import io308    from pypdf import PdfReader, PdfWriter309 310    base = PdfReader(io.BytesIO(base_pdf))311    overlay_reader = PdfReader(io.BytesIO(overlay_pdf))312    overlay_page = overlay_reader.pages[0]313    writer = PdfWriter()314 315    total = len(base.pages)316    first_content = skip_front            # e.g. page index 2317    last_content = total - skip_back - 1  # e.g. total-5318 319    for i, pg in enumerate(base.pages):320        if first_content <= i <= last_content:321            pg.merge_page(overlay_page)322        writer.add_page(pg)323 324    buf = io.BytesIO()325    writer.write(buf)326    return buf.getvalue()327