1"""Website screenshots and PDFs with Playwright (Chromium)."""
2from __future__ import annotations
3
4import asyncio
5import io
6import urllib.parse
7from typing import Awaitable, Callable
8
9from PIL import Image
10from playwright.async_api import Browser, Error as PlaywrightError, Route
11
12DEVICES = {
13 'desktop': {'viewport': {'width': 1920, 'height': 1080}},
14 'laptop': {'viewport': {'width': 1366, 'height': 768}},
15 'tablet': 'iPad Pro 11',
16 'mobile': 'iPhone 15 Pro',
17 'android': 'Galaxy S24',
18}
19
20HIDE_BANNERS_JS = r"""
21() => {
22 const re = /(cookie|consent|gdpr|cmp|onetrust|cookiebot|didomi|usercentrics|privacy-banner|cc-window|qc-cmp)/i;
23 let n = 0;
24 for (const el of document.querySelectorAll('body *')) {
25 const s = getComputedStyle(el);
26 if (s.position !== 'fixed' && s.position !== 'sticky') continue;
27 const tag = `${el.id} ${typeof el.className === 'string' ? el.className : ''} ${el.getAttribute('aria-label') || ''}`;
28 if (re.test(tag)) { el.style.setProperty('display', 'none', 'important'); n++; }
29 }
30 document.documentElement.style.overflow = 'auto';
31 document.body && (document.body.style.overflow = 'auto');
32 return n;
33}
34"""
35SCROLL_JS = r"""
36async (maxHeight) => {
37 const step = Math.max(window.innerHeight * 0.8, 400);
38 for (let y = 0; y < Math.min(document.body.scrollHeight, maxHeight); y += step) {
39 window.scrollTo(0, y);
40 await new Promise(r => setTimeout(r, 150));
41 }
42 window.scrollTo(0, 0);
43 await new Promise(r => setTimeout(r, 300));
44}
45"""
46MAX_FULL_HEIGHT = 20000
47
48
49class ShotError(Exception):
50 pass
51
52
53def device_options(playwright, device: str, width: int, height: int, dark: bool) -> dict:
54 spec = DEVICES.get(device, DEVICES['desktop'])
55 opts = dict(playwright.devices[spec]) if isinstance(spec, str) else dict(spec)
56 opts.pop('default_browser_type', None)
57 if device in ('desktop', 'laptop') and width:
58 opts['viewport'] = {'width': width, 'height': height or opts['viewport']['height']}
59 opts['color_scheme'] = 'dark' if dark else 'light'
60 opts['ignore_https_errors'] = True
61 return opts
62
63
64async def capture(browser: Browser, ctx_opts: dict, target: dict, opts: dict,
65 host_ok: Callable[[str], Awaitable[bool]]) -> dict:
66 """target: {'url': ...} or {'html': ...}. Returns {'data': bytes, 'contentType', 'title', 'finalUrl', 'status',
67 'width', 'height', 'bannersHidden'}."""
68 context = await browser.new_context(**ctx_opts)
69 blocked = []
70
71 async def guard(route: Route) -> None:
72 url = route.request.url
73 scheme = urllib.parse.urlparse(url).scheme
74 if scheme in ('data', 'blob', 'about'):
75 return await route.continue_()
76 host = urllib.parse.urlparse(url).hostname or ''
77 if scheme not in ('http', 'https') or not await host_ok(host):
78 blocked.append(url[:200])
79 return await route.abort('blockedbyclient')
80 await route.continue_()
81
82 await context.route('**/*', guard)
83 page = await context.new_page()
84 page.set_default_timeout(opts['timeoutSecs'] * 1000)
85 try:
86 status = None
87 if 'url' in target:
88 try:
89 resp = await page.goto(target['url'], wait_until=opts['waitUntil'])
90 except PlaywrightError as e:
91 if blocked and blocked[0].split('#')[0] == target['url'].split('#')[0]:
92 raise ShotError('the URL points to a private address') from e
93 if 'Timeout' in str(e) and opts['waitUntil'] == 'networkidle':
94 resp = None
95 else:
96 raise ShotError('page did not load: ' + str(e).splitlines()[0][:200]) from e
97 status = resp.status if resp else None
98 else:
99 await page.set_content(target['html'], wait_until='load')
100 if opts['scroll']:
101 await page.evaluate(SCROLL_JS, MAX_FULL_HEIGHT)
102 hidden = await page.evaluate(HIDE_BANNERS_JS) if opts['hideBanners'] else 0
103 if opts['css']:
104 await page.add_style_tag(content=opts['css'])
105 if opts['delayMs']:
106 await asyncio.sleep(opts['delayMs'] / 1000)
107 fmt = opts['format']
108 if fmt == 'pdf':
109 await page.emulate_media(media='screen' if opts['pdfScreenMedia'] else 'print')
110 margin = opts['pdfMargin']
111 data = await page.pdf(format=opts['pdfFormat'], landscape=opts['pdfLandscape'], print_background=True,
112 margin={k: margin for k in ('top', 'right', 'bottom', 'left')})
113 ctype = 'application/pdf'
114 else:
115 shot_fmt = 'jpeg' if fmt == 'jpeg' else 'png'
116 kw = {'type': shot_fmt}
117 if shot_fmt == 'jpeg':
118 kw['quality'] = opts['quality']
119 if opts['selector']:
120 el = await page.query_selector(opts['selector'])
121 if not el:
122 raise ShotError(f'selector "{opts["selector"]}" was not found on the page')
123 data = await el.screenshot(**kw)
124 else:
125 full = opts['fullPage']
126 if full:
127 h = await page.evaluate('() => document.documentElement.scrollHeight')
128 if h > MAX_FULL_HEIGHT:
129 vw = page.viewport_size['width']
130 kw['clip'] = {'x': 0, 'y': 0, 'width': vw, 'height': MAX_FULL_HEIGHT}
131 data = await page.screenshot(full_page=full, **kw)
132 ctype = f'image/{shot_fmt}'
133 if fmt == 'webp':
134 buf = io.BytesIO()
135 Image.open(io.BytesIO(data)).save(buf, 'WEBP', quality=opts['quality'])
136 data, ctype = buf.getvalue(), 'image/webp'
137 size = None
138 if fmt != 'pdf':
139 with Image.open(io.BytesIO(data)) as im:
140 size = im.size
141 return {'data': data, 'contentType': ctype, 'title': await page.title(), 'finalUrl': page.url,
142 'status': status, 'width': size[0] if size else None, 'height': size[1] if size else None,
143 'bannersHidden': hidden, 'blockedRequests': len(blocked)}
144 except PlaywrightError as e:
145 raise ShotError(str(e).splitlines()[0][:300]) from e
146 finally:
147 await context.close()
