From 763becd8a40bff92c0ff94f4430461f12df52029 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Corentin=20S=C3=A9chet?= Date: Mon, 11 Apr 2022 12:22:33 +0200 Subject: [PATCH] common: use a wrapper for browser context, as route on the context raise exception but on page not --- frontools/cache.py | 8 +++--- frontools/screenshot.py | 54 +++++++++++++++++++---------------------- frontools/sources.py | 33 ++++++++++++++++++------- 3 files changed, 54 insertions(+), 41 deletions(-) diff --git a/frontools/cache.py b/frontools/cache.py index 2956404..0121435 100644 --- a/frontools/cache.py +++ b/frontools/cache.py @@ -79,11 +79,13 @@ class FileCache(Cache[ResourceType]): def _get_cache_file_path(self, key: str) -> Path: key_slug = _get_key_slug(key) cache_directory = xdg_cache_home() / "frontools" / self._name + file_path = cache_directory.joinpath(*key_slug.split("&")) + file_directory = file_path.parent - if not cache_directory.is_dir(): - cache_directory.mkdir(parents=True) + if not file_directory.is_dir(): + file_directory.mkdir(parents=True) - return cache_directory / key_slug + return file_path @abstractmethod def _serialize(self, resource: ResourceType) -> bytes: diff --git a/frontools/screenshot.py b/frontools/screenshot.py index d29436a..439fb9a 100644 --- a/frontools/screenshot.py +++ b/frontools/screenshot.py @@ -7,6 +7,7 @@ from PIL import Image, ImageChops from playwright.async_api import BrowserContext from frontools import Config +from frontools.sources import Browser from frontools.utils import ( get_default_screenshot_directory, get_url_slug, @@ -55,43 +56,38 @@ async def screenshot_diff( ) for (theme, url) in urls ], - nb_workers=2, + nb_workers=3, ) -async def _diff_url( - left: BrowserContext, right: BrowserContext, url: str, output_path: Path -) -> None: - left_bytes = await _screenshot_url(left, url) - right_bytes = await _screenshot_url(right, url) +async def _diff_url(left: Browser, right: Browser, url: str, output_path: Path) -> None: + try: + 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) + with NamedTemporaryFile(mode="wb") as left_file: + left_file.write(left_bytes) + left_image = Image.open(left_file.name) - with NamedTemporaryFile(mode='wb') as right_file: - right_file.write(right_bytes) - right_image = Image.open(right_file.name) + with NamedTemporaryFile(mode="wb") as right_file: + right_file.write(right_bytes) + right_image = Image.open(right_file.name) - diff = ImageChops.difference(left_image, right_image) + diff = ImageChops.difference(left_image, right_image) - if not diff.getbbox(): # images are the same - return + url_slug = get_url_slug(url) + if not output_path.is_dir(): + output_path.mkdir() - url_slug = get_url_slug(url) - if not output_path.is_dir(): - output_path.mkdir() + with open(output_path / f"{url_slug}_left", "wb") as screenshot_file: + screenshot_file.write(left_bytes) - with open(output_path / f"{url_slug}_left", "wb") as screenshot_file: - screenshot_file.write(left_bytes) - - with open(output_path / f"{url_slug}_right", "wb") as screenshot_file: - screenshot_file.write(right_bytes) + with open(output_path / f"{url_slug}_right", "wb") as screenshot_file: + screenshot_file.write(right_bytes) + except Exception as ex: + print(f"Error while diffing {url} : {ex}") -async def _screenshot_url(browser: BrowserContext, url: str) -> bytes: - page = await browser.new_page() - await page.goto(url) - await page.wait_for_load_state("networkidle") - - return await page.screenshot(full_page=True) +async def _screenshot_url(browser: Browser, url: str) -> bytes: + async with browser.load_page(url) as page: + return await page.screenshot(full_page=True) diff --git a/frontools/sources.py b/frontools/sources.py index 018eb18..c434ded 100644 --- a/frontools/sources.py +++ b/frontools/sources.py @@ -6,11 +6,27 @@ from re import compile as re_compile from typing import AsyncGenerator, Optional, cast from aiohttp import ClientSession -from playwright.async_api import BrowserContext, Route, ViewportSize, async_playwright +from playwright.async_api import BrowserContext, Route, ViewportSize, async_playwright, Page from frontools.cache import Cache +class Browser: + def __init__(self, source: 'Source', browser_context: BrowserContext) -> None: + """Wraps a browser instance, with helpers methods to load pages.""" + self._source = source + self._browser_context = browser_context + + @asynccontextmanager + async def load_page(self, url: str) -> AsyncGenerator[Page, None]: + page = await self._browser_context.new_page() + await page.route("*", self._source.route) + await page.goto(url) + await page.wait_for_load_state("networkidle") + yield page + await page.close() + + class Source(ABC): """Base class for sources""" @@ -21,7 +37,7 @@ class Source(ABC): @asynccontextmanager async def get_browser( self, width: Optional[int] = None, height: Optional[int] = None - ) -> AsyncGenerator[BrowserContext, None]: + ) -> AsyncGenerator[Browser, None]: """Return a Playwright browser that will eventually get files from local cache""" viewport: ViewportSize = cast( @@ -32,17 +48,16 @@ class Source(ABC): assert height is not None viewport = dict(width=width, height=height) - async def _cache_route(route: Route) -> None: - content = await self.get_url(route.request.url) - route.fulfill(body=content) - async with async_playwright() as pwright: - browser = await pwright.firefox.launch() + browser = await pwright.firefox.launch(headless=True) context = await browser.new_context(viewport=viewport) - await context.route("*", _cache_route) - yield context + yield Browser(self, context) await browser.close() + async def route(self, route: Route) -> None: + content = await self.get_url(route.request.url) + await route.fulfill(body=content, status=200) + class CachedSource(Source): """Source loading urls from the internet."""