misc-csechet/frontools/screenshot.py

130 lines
3.9 KiB
Python

"""Pages screenshots"""
from datetime import datetime
from os.path import expandvars
from pathlib import Path
from re import compile as re_compile
from tempfile import NamedTemporaryFile
from typing import Optional
from PIL import Image, ImageChops
from playwright.async_api import Error as PlaywrightError
from xdg import xdg_config_home
from frontools.config import Config
from frontools.sources import Browser
from frontools.utils import get_url_slug, report_progress
async def screenshot_diff(
config: Config,
right_source_name: str,
output_directory: Optional[str],
continue_: bool = False,
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, exist_ok=True)
left_source = config.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:
await report_progress(
"Screenshoting",
[
(
url,
_diff_url(
left_browser,
right_browser,
url,
output_path,
theme_name,
continue_
),
)
for (theme_name, url) in config.urls
],
nb_workers=3,
)
async def _diff_url(
left: Browser,
right: Browser,
url: str,
output_path: Path,
theme_name: str,
continue_: bool
) -> None:
url_slug = get_url_slug(url)
left_filename = output_path / f"{theme_name}_{url_slug}_left"
right_filename = output_path / f"{theme_name}_{url_slug}_right"
if left_filename.is_file() and right_filename.is_file() and continue_:
return
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
output_path.mkdir(parents=True, exist_ok=True)
with open(left_filename, "wb") as screenshot_file:
screenshot_file.write(left_bytes)
with open(right_filename, "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()
def _get_default_screenshot_directory() -> Path:
pictures_dir_re = re_compile(r'\w*XDG_PICTURES_DIR\w*\="(?P<path>.*)"\w*$')
pictures_dir = None
with open(
Path(xdg_config_home()) / "user-dirs.dirs", encoding="utf-8"
) as user_dirs:
for line in user_dirs.readlines():
match = pictures_dir_re.match(line)
if match is None:
continue
path_string = match.group("path")
path_string = expandvars(path_string)
pictures_dir = Path(path_string)
break
if pictures_dir is None:
pictures_dir = Path.home()
return pictures_dir / datetime.now().strftime("%Y-%m-%d - %H-%M-%S")