screenshot: taking screenshots from one source at a time

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_
     )
 
 

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"""

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: