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

app/core/service/playwright/playwright_context.py

@@ -1,139 +1,88 @@
 from __future__ import annotations
 
-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: ClassVar[str] = (
-        "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()
-        
+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"
+