Update app/core/service/playwright/playwright_context.py

app/core/service/playwright/playwright_context.py

@@ -1,88 +1,139 @@
 from __future__ import annotations
 
-from typing import Literal
-
-from pydantic import BaseModel, Field, HttpUrl
-
-
-class ViewPortModel(BaseModel):
-    """Page viewport
-
-    Attributes:
-        width (int):
-            viewport width.
-
-        height (int):
-            viewport height.
-    """
-
-    width: int = 1280
-    height: int = 720
-
-
-class PageModel(BaseModel):
-    """Page attrs
-
-    Attributes:
-        color_scheme (Literal["light", "dark", "no-preference"] | None):
-            Page color.
-
-        java_script_enabled (bool | None):
-            Whether or not to enable JavaScript in the context. Defaults to true.
-
-        viewport (ViewPortModel | None):
-            Sets a consistent viewport for each page. Defaults to an 1280x720 viewport.
-
-        no_viewport (bool | None):
-            Does not enforce fixed viewport, allows resizing window in the headed mode.
-
-        proxy (dict | None):
-            Proxy to be used for all requests. HTTP and SOCKS proxies are supported. Example: proxy={'server': 'http://proxy.example.com:3128'}
-    """  # noqa: E501
-
-    color_scheme: Literal["light", "dark", "no-preference"] | None = "no-preference"
-    java_script_enabled: bool | None = True
-    viewport: ViewPortModel | None = None
-    no_viewport: bool | None = False
-    proxy: dict | None = None
-
-
-class GetContentModel(BaseModel):
-    """Webpage to request and parse.
-
-    Attributes:
-        url (HttpUrl):
-            Url to request.
-
-        new_browser (bool | None):
-            Whether you want to make a new browser context or not.
-
-        query_selector (str | None):
-            Used to locate a selector.
-
-        ms_delay (int):
-            A delay before performing a task after requesting the url.
-    """
-
-    url: HttpUrl
-    new_browser: bool | None = False
-    query_selector: str | None = None
-    ms_delay: int = Field(default=0.0, ge=0, le=15_000)
-
-
-class ScreenshotModel(GetContentModel):
-    """Screenshot schemas
-
-    Attributes:
-
-        full_page (bool | None): Whether you want a full page screenshot or not.
-
-        image_type (Literal["png", "jpeg"]):
-            The image type of screenshot.
-    """
-
-    full_page: bool | None = Field(
-        default = False,
-        description = "Whether you want a full page screenshot or not.")
-    image_type: Literal["png", "jpeg"] = "jpeg"
-    
+from typing import AsyncIterator, Awaitable, ClassVar
+
+from playwright.async_api import (  # noqa: F401
+    Browser,
+    BrowserContext,
+    Page,
+    TimeoutError,
+    async_playwright,
+)
+
+from .models import GetContentModel, PageModel, ScreenshotModel  # noqa: TCH001
+
+
+class AsyncMixin:
+    """Experimental: making awaitable class."""
+
+    async def __ainit__(self) -> None:
+        """Initialize the class."""
+
+    def __await__(self) -> AsyncIterator[Awaitable]:
+        """Make the class awaitable."""
+        return self.__ainit__().__await__()
+
+
+class PlaywrightInstance(AsyncMixin):
+    """This class is designed to keep playwright browser instance open for reusability and scalability handling api requests."""  # noqa: E501
+
+    HEADERS = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36"  # noqa: E501
+
+    FIREFOX_USER_PREFS: ClassVar[dict[str, int | str]] = {
+        "extensions.enabledScopes": 1,
+        "extensions.autoDisableScopes": 1,
+        "dom.webdriver.enabled": False,
+        "useAutomationExtension": False,
+        "general.useragent.override": HEADERS,
+    }
+
+    def __init__(self) -> None:
+        """Initialize the class."""
+        self.playwright: async_playwright | None = None
+        self.browser: Browser | None = None
+
+    async def __ainit__(self) -> PlaywrightInstance:
+        if not self.playwright:
+            self.playwright = await async_playwright().start()
+            self.browser = await self.playwright.firefox.launch(
+                firefox_user_prefs=self.FIREFOX_USER_PREFS,
+            )
+        return self
+
+    async def new_context_page(
+        self,
+        browser: Browser,
+        screenshot_model: GetContentModel,
+        page_model: PageModel,
+    ) -> tuple[BrowserContext | None, Page]:
+        """create a brwoser or new browser context page.
+
+        Parameters:
+            browser (Browser):
+                The Playwright Browser instance to create a new context in.
+            screenshot_model (GetContentModel):
+                A pydantic BaseModel instance containing the configuration for the screenshot.
+            page_model (PageModel):
+                A pydantic BaseModel instance containing the configuration for the page.
+
+        Returns:
+            tuple: BrowserContext and Page
+        """
+        params = {
+            "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,
+        }
+
+        if not screenshot_model.new_browser:
+            return None, await browser.new_page(**params)
+
+        new_context = await browser.new_context(**params)
+        return new_context, await new_context.new_page()
+
+    async def screenshot(
+        self,
+        screenshot_model: ScreenshotModel,
+        page_model: PageModel,
+    ) -> bytes:
+        """Take a screenshot of a webpage url.
+
+        Parameters:
+            screenshot_model (ScreenshotModel):
+                A pydantic BaseModel instance containing the configuration for the screenshot.
+            page_model (PageModel):
+                A pydantic BaseModel instance containing the configuration for the page.
+
+        Returns:
+            bytes: The screenshot data in bytes.
+        """
+        context, page = await self.new_context_page(
+            screenshot_model=screenshot_model,
+            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:
+            screenshot_data: bytes = await screenshot_locator.screenshot(
+                type=screenshot_model.image_type,
+            )
+        else:
+            screenshot_data: bytes = await page.screenshot(
+                full_page=screenshot_model.full_page,
+                type=screenshot_model.image_type,
+            )
+
+        await page.close()
+
+        if context:
+            await context.close()
+
+        return screenshot_data
+
+    async def close_instance(self) -> None:
+        """For manual closing of playwright if needed."""
+        if self.playwright:
+            await self.browser.close()
+            await self.playwright.stop()
+