misc-csechet/frontools/screenshot.py

105 lines
3.0 KiB
Python

"""Pages screenshots"""
from pathlib import Path
from tempfile import NamedTemporaryFile
from typing import Optional
from PIL import Image, ImageChops
from playwright.async_api import Error as PlaywrightError
from frontools.config import Config
from frontools.sources import Browser
from frontools.utils import (
get_default_screenshot_directory,
get_url_slug,
report_progress,
)
async def screenshot_diff(
config: Config,
right_source_name: str,
output_directory: Optional[str],
screen_width: Optional[int] = None,
) -> None:
"""Compare pages with or without local css"""
if output_directory is None:
output_path = get_default_screenshot_directory()
else:
output_path = Path(output_directory)
output_path.mkdir(parents=True)
left_source = config.default_source
right_source = config.get_source(right_source_name)
async with left_source.get_browser(width=screen_width) as left_browser:
async with right_source.get_browser(width=screen_width) as right_browser:
urls = [
(site_name, url)
for (site_name, site) in config.sites
for url in site.urls
]
await report_progress(
"Screenshoting",
[
(
url,
_diff_url(
left_browser,
right_browser,
url,
output_path,
site_name,
),
)
for (site_name, url) in urls
],
nb_workers=3,
)
async def _diff_url(
left: Browser,
right: Browser,
url: str,
output_path: Path,
site_name: str,
) -> None:
left_bytes = await _screenshot_url(left, url)
right_bytes = await _screenshot_url(right, url)
with NamedTemporaryFile(mode="wb") as left_file:
left_file.write(left_bytes)
left_image = Image.open(left_file.name).convert("RGB")
with NamedTemporaryFile(mode="wb") as right_file:
right_file.write(right_bytes)
right_image = Image.open(right_file.name).convert("RGB")
diff = ImageChops.difference(left_image, right_image)
if not diff.getbbox():
return
url_slug = get_url_slug(url)
if not output_path.is_dir():
output_path.mkdir()
with open(output_path / f"{site_name}_{url_slug}_left", "wb") as screenshot_file:
screenshot_file.write(left_bytes)
with open(output_path / f"{site_name}_{url_slug}_right", "wb") as screenshot_file:
screenshot_file.write(right_bytes)
async def _screenshot_url(browser: Browser, url: str) -> bytes:
async with browser.load_page(url) as page:
try:
return await page.screenshot(full_page=True)
except PlaywrightError:
pass
# Exception raisen when taking a screenshot of a to large page, retry without full_page
return await page.screenshot()