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

app/core/service/playwright/playwright_context.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, AsyncContextManager
+from typing import AsyncIterator, Awaitable, ClassVar
 
 from playwright.async_api import (  # noqa: F401
     Browser,
@@ -10,66 +10,86 @@ from playwright.async_api import (  # noqa: F401
     async_playwright,
 )
 
-from .models import GetContentModel, PageModel, ScreenshotModel
+from .models import GetContentModel, PageModel, ScreenshotModel  # noqa: TCH001
 
-if TYPE_CHECKING:
-    from types import TracebackType
 
+class AsyncMixin:
+    """experimental: making awaitable class."""
 
-class AsyncPlaywrightContext:
-    HEADERS = "Mozilla/5.0 (Macintosh; Intel Mac OS X 14_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36"  # noqa: E501
+    async def ainit(self) -> None:
+        pass
+
+    def __await__(self) -> AsyncIterator[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: 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:
-        self.playwright = None
-        self.browser = None
-        self.default_context = None
+        self.playwright: async_playwright = None
+        self.browser: Browser = None
 
-    async def __aenter__(self) -> AsyncContextManager:
+    async def ainit(self) -> PlaywrightInstance:
+        """starts playwright and browser instance."""
         if not self.playwright:
             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": self.HEADERS,
-                },
+                firefox_user_prefs=self.FIREFOX_USER_PREFS,
             )
-            self.default_context = await self.browser.new_context()
         return self
 
     async def new_context_page(
         self,
-        screenshot_model: GetContentModel,
         browser: Browser,
+        screenshot_model: GetContentModel,
         page_model: PageModel,
-    ) -> Page:
-            
+    ) -> tuple[BrowserContext, 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,
+            "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 await self.browser.new_page(**params)
-        
-        new_context = await browser.new_context(**params)
-        return await new_context.new_page()
+            return None, await self.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:
-        page = await self.new_context_page(
-            screenshot_model = screenshot_model,
-            browser=self.browser,
-            page_model=page_model)
+        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))
 
@@ -82,55 +102,20 @@ class AsyncPlaywrightContext:
         )
 
         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,
-        page_model: PageModel,
-    ) -> str:
-        page = await self.new_context_page(
-            screenshot_model = GetContentModel,
-            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 context:
+            await context.close()
 
-        if wait_locator:
-            await wait_locator.wait_for()
-        html = page.content()
-        await page.close()
-        return html
+        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()
-            self.browser = None
-            self.playwright = None
-
-    async def __aexit__(
-        self,
-        typ: type[BaseException] | None,
-        exc: BaseException | None,
-        tb: TracebackType | None,
-    ) -> None:
-        print('speed test')
-        # if self.browser:
-        #     await self.browser.close()
-        # if self.playwright:
-        #     await self.playwright.stop()
+