Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| from typing import TYPE_CHECKING, AsyncContextManager | |
| from playwright.async_api import ( # noqa: F401 | |
| Browser, | |
| BrowserContext, | |
| Page, | |
| TimeoutError, | |
| async_playwright, | |
| ) | |
| from .models import GetContentModel, PageModel, ScreenshotModel # noqa: TCH001 | |
| if TYPE_CHECKING: | |
| from types import TracebackType | |
| class AsyncPlaywrightContext: | |
| async def __aenter__(self) -> AsyncContextManager: | |
| self.playwright = await async_playwright().start() | |
| self.browser = await self.playwright.firefox.launch( | |
| firefox_user_prefs={ | |
| "extensions.enabledScopes": 1, | |
| "extensions.autoDisableScopes": 1, | |
| "dom.webdriver.enabled": False, | |
| "useAutomationExtension": False, | |
| "general.useragent.override": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/83.0.4103.61 Safari/537.36", # noqa: E501 | |
| }, | |
| ) | |
| return self | |
| async def new_browser_page(self, browser: Browser, page_model: PageModel) -> Page: | |
| return await browser.new_page( | |
| color_scheme=page_model.color_scheme, | |
| java_script_enabled=page_model.java_script_enabled, | |
| no_viewport=page_model.no_viewport, | |
| proxy=page_model.proxy.model_dump() if page_model.proxy else None, | |
| viewport=page_model.viewport.model_dump() if page_model.viewport else None, | |
| ) | |
| async def screenshot( | |
| self, | |
| screenshot_model: ScreenshotModel, | |
| page_model: PageModel, | |
| ) -> bytes: | |
| page = await self.new_browser_page(browser=self.browser, page_model=page_model) | |
| await page.goto(str(screenshot_model.url)) | |
| await page.wait_for_timeout(screenshot_model.ms_delay) | |
| screenshot_locator = ( | |
| page.locator(screenshot_model.query_selector) | |
| if screenshot_model.query_selector | |
| else None | |
| ) | |
| if screenshot_locator: | |
| if screenshot_model.wait_selector: | |
| await screenshot_locator.wait_for() | |
| screenshot_data: bytes = await screenshot_locator.screenshot() | |
| else: | |
| screenshot_data: bytes = await page.screenshot(full_page=screenshot_model.full_page) | |
| await page.close() | |
| return screenshot_data | |
| async def get_content( | |
| self, | |
| get_content_model: GetContentModel, | |
| ) -> str: | |
| page = await self.new_browser_page(browser=self.browser, page_model=page_model) | |
| await page.goto(str(get_content_model.url)) | |
| await page.wait_for_timeout(get_content_model.ms_delay) | |
| wait_locator = ( | |
| await page.locator(get_content_model.query_selector) | |
| if get_content_model.query_selector | |
| else None | |
| ) | |
| if wait_locator: | |
| await wait_locator.wait_for() | |
| html = page.content() | |
| await page.close() | |
| return html | |
| async def __aexit__( | |
| self, | |
| typ: type[BaseException] | None, | |
| exc: BaseException | None, | |
| tb: TracebackType | None, | |
| ) -> None: | |
| if self.browser: | |
| await self.browser.close() | |
| if self.playwright: | |
| await self.playwright.stop() | |