Initial commit

Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.11 AS builder
+ 
+WORKDIR /app
+ 
+RUN python3 -m venv venv
+ENV VIRTUAL_ENV=/app/venv
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+ 
+COPY requirements.txt .
+RUN pip install -r requirements.txt
+ 
+# Stage 2
+FROM python:3.11 AS runner
+ 
+WORKDIR /app
+ 
+COPY --from=builder /app/venv venv
+ 
+ENV VIRTUAL_ENV=/app/venv
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+ 
+EXPOSE 8000
+ 
+CMD [ "python screenshot/main.py" ]

pyproject.toml

@@ -0,0 +1,18 @@
+[tool.ruff]
+line-length = 100
+
+[tool.ruff.lint]
+select = ["ALL"]
+ignore = [
+  "CPY001", # copyright above code
+  "D", # sphinx not support 
+  ]
+
+[tool.mypy]
+disallow_untyped_defs = true
+show_error_codes = true
+no_implicit_optional = true
+warn_return_any = true
+warn_unused_ignores = true
+exclude = ["tests"]
+python_version = "3.10"

screenshot/main.py

@@ -0,0 +1,10 @@
+import uvicorn
+from fastapi import FastAPI
+from routers.screenshot import routers as screenshot_router
+
+app = FastAPI()
+
+app.include_router(screenshot_router)
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)  # noqa: S104

screenshot/routers/screenshot.py

@@ -0,0 +1,93 @@
+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")