diff --git a/frontools/cli.py b/frontools/cli.py index 27b254c..18b09d4 100644 --- a/frontools/cli.py +++ b/frontools/cli.py @@ -10,7 +10,7 @@ from click import Path as PathArgument from click import argument, group, option, pass_context, pass_obj from frontools.config import Config -from frontools.screenshot import screenshot_diff +from frontools.screenshot import screenshot from frontools.sources import CachedSource @@ -119,8 +119,7 @@ def prune_caches(cache_names: list[str]) -> None: CachedSource.prune(cache_names) -@main.command(name="screenshot-diff") -@argument("source", type=str) +@main.command(name="screenshot") @option("-o", "--output-directory", type=PathArgument(), default=None) @option("-w", "--screen-width", type=int, default=None) @option( @@ -133,16 +132,15 @@ def prune_caches(cache_names: list[str]) -> None: ) @_async_command @pass_obj -async def screenshot_diff_cli( +async def screenshot_cli( config: Config, - source: str, output_directory: Optional[str], screen_width: Optional[int], continue_: bool, ) -> None: """Generate screenshot diffs""" - await screenshot_diff( - config, source, output_directory, screen_width=screen_width, continue_=continue_ + await screenshot( + config, output_directory, screen_width=screen_width, continue_=continue_ ) diff --git a/frontools/config.py b/frontools/config.py index d8d3712..322214f 100644 --- a/frontools/config.py +++ b/frontools/config.py @@ -80,6 +80,11 @@ class Config: """get the default source for this context""" return self.get_source(self._source_name) + @property + def source_name(self) -> str: + """get the default source name for this config""" + return self._source_name + @property def urls(self) -> Iterable[tuple[str, str]]: """Return themes configured for this context""" diff --git a/frontools/screenshot.py b/frontools/screenshot.py index 7f2bada..b69356a 100644 --- a/frontools/screenshot.py +++ b/frontools/screenshot.py @@ -3,10 +3,8 @@ 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 @@ -15,9 +13,8 @@ from frontools.sources import Browser from frontools.utils import get_url_slug, report_progress -async def screenshot_diff( +async def screenshot( config: Config, - right_source_name: str, output_directory: Optional[str], continue_: bool = False, screen_width: Optional[int] = None, @@ -30,81 +27,46 @@ async def screenshot_diff( 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) + source = config.source - 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 with source.get_browser(width=screen_width) as browser: + await report_progress( + "Screenshoting", + [ + ( + url, + _screenshot_url(browser, output_path, url, theme_name, config.source_name, continue_), + ) + for (theme_name, url) in config.urls + ], + nb_workers=3, + ) -async def _diff_url( - left: Browser, - right: Browser, - url: str, +async def _screenshot_url( + browser: Browser, output_path: Path, + url: str, theme_name: str, - continue_: bool + source_name: str, + continue_: bool, ) -> None: url_slug = get_url_slug(url) - left_filename = output_path / f"{theme_name}_{url_slug}_left.png" - right_filename = output_path / f"{theme_name}_{url_slug}_right.png" + filename = output_path / f"{theme_name}_{url_slug}_{source_name}.png" - if left_filename.is_file() and right_filename.is_file() and continue_: + if 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) + screenshot = await page.screenshot(full_page=True) except PlaywrightError: - pass + # Exception raisen when taking a screenshot of a to large page, retry without full_page + screenshot = await page.screenshot() - # Exception raisen when taking a screenshot of a to large page, retry without full_page - return await page.screenshot() + with open(filename, "wb") as screenshot_file: + screenshot_file.write(screenshot) def _get_default_screenshot_directory() -> Path: