misc-csechet/frontools/screenshot.py

107 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 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],
resolution: Optional[str] = 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)
width: Optional[int] = None
height: Optional[int] = None
if resolution is not None:
[width_str, height_str] = resolution.split("x")
width = int(width_str)
height = int(height_str)
async with left_source.get_browser(width=width, height=height) as left_browser:
async with right_source.get_browser(
width=width, height=height
) 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:
return await page.screenshot(full_page=True)