screenshot/routers/screenshot.py

from __future__ import annotations

import logging
from typing import IO, TYPE_CHECKING, AsyncContextManager, Literal

from fastapi import APIRouter, HTTPException
from fastapi.responses import Response
from playwright.async_api import BrowserContext, TimeoutError, async_playwright
from pydantic import BaseModel, HttpUrl

if TYPE_CHECKING:
    from types import TracebackType


router = APIRouter()


class ViewPort(BaseModel):
    width: int = 1280
    height: int = 720


class ScreenshotItems(BaseModel):
    url: HttpUrl
    full_page: bool | None = False
    query_selector: str | None = None

    viewport: ViewPort | None = None
    color_scheme: Literal["light", "dark", "no-preference"] | None = "no-preference"
    bypass_csp: bool | None = False
    java_script_enabled: bool | None = True
    proxy: dict | None = None
    is_mobile: bool | None = False
    no_viewport: bool | None = False


class ScreenShot:
    async def __aenter__(self) -> AsyncContextManager[ScreenShot]:
        self.playwright = await async_playwright().start()
        self.browser = await self.playwright.chromium.launch(
            args=["--disable-extensions"],
            chromium_sandbox=True,
        )
        return self

    async def browser_context(self, items: ScreenshotItems) -> BrowserContext:
        return await self.browser.new_context(
            viewport=items.viewport.model_dump() if items.viewport else None,
            color_scheme=items.color_scheme,
            bypass_csp=items.bypass_csp,
            java_script_enabled=items.java_script_enabled,
            proxy=items.proxy.model_dump() if items.proxy else None,
            is_mobile=items.is_mobile,
            no_viewport=items.no_viewport,
        )

    async def capture(self, items: ScreenshotItems) -> IO[bytes]:
        context: BrowserContext = await self.browser_context(items)
        page = await context.new_page()
        await page.goto(str(items.url))

        if items.query_selector:
            page = page.locator(items.query_selector)

        screenshot_data = await page.screenshot(full_page=items.full_page)
        await context.close()
        return screenshot_data

    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()


@router.post("/screenshot")
async def screenshot(data: ScreenshotItems) -> Response:
    async with ScreenShot() as sc:
        try:
            response = await sc.capture(items=data)
            return Response(content=response, media_type="image/png")
        except TimeoutError as e:
            raise HTTPException(
                status_code=504,
                detail=f"An error occurred while generating the screenshot: {e}",
            ) from e
        except Exception:
            logging.exception("screenshot unhandled error")