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

app/core/service/playwright/playwright_context.py

@@ -14,19 +14,22 @@ from .models import GetContentModel, PageModel, ScreenshotModel  # noqa: TCH001
 
 
 class AsyncMixin:
-    """experimental: making awaitable class."""
+    """Experimental: making awaitable class."""
 
-    async def ainit(self) -> None:
-        pass
+    async def __ainit__(self) -> None:
+        """Initialize the class."""
 
     def __await__(self) -> AsyncIterator[Awaitable]:
-        return self.ainit().__await__()
+        """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: 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
+    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,
@@ -37,11 +40,11 @@ class PlaywrightInstance(AsyncMixin):
     }
 
     def __init__(self) -> None:
-        self.playwright: async_playwright = None
-        self.browser: Browser = None
+        """Initialize the class."""
+        self.playwright: async_playwright | None = None
+        self.browser: Browser | None = None
 
-    async def ainit(self) -> PlaywrightInstance:
-        """starts playwright and browser instance."""
+    async def __ainit__(self) -> PlaywrightInstance:
         if not self.playwright:
             self.playwright = await async_playwright().start()
             self.browser = await self.playwright.firefox.launch(
@@ -54,7 +57,7 @@ class PlaywrightInstance(AsyncMixin):
         browser: Browser,
         screenshot_model: GetContentModel,
         page_model: PageModel,
-    ) -> tuple[BrowserContext, Page]:
+    ) -> tuple[BrowserContext | None, Page]:
         """create a brwoser or new browser context page.
 
         Parameters:
@@ -73,11 +76,11 @@ class PlaywrightInstance(AsyncMixin):
             "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,
+            "viewport": page_model.viewport,
         }
 
         if not screenshot_model.new_browser:
-            return None, await self.browser.new_page(**params)
+            return None, await browser.new_page(**params)
 
         new_context = await browser.new_context(**params)
         return new_context, await new_context.new_page()
@@ -87,24 +90,39 @@ class PlaywrightInstance(AsyncMixin):
         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,
+            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
+            else None,
         )
 
         if screenshot_locator:
-            screenshot_data: bytes = await screenshot_locator.screenshot()
+            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)
+            screenshot_data: bytes = await page.screenshot(
+                full_page=screenshot_model.full_page, type=screenshot_model.image_type
+            )
 
         await page.close()
 
@@ -114,7 +132,7 @@ class PlaywrightInstance(AsyncMixin):
         return screenshot_data
 
     async def close_instance(self) -> None:
-        """for manual closing of playwright if needed"""
+        """For manual closing of playwright if needed."""
         if self.playwright:
             await self.browser.close()
             await self.playwright.stop()