version 1.0

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,132 @@
+# requirements dev
+.ruff_cache
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

Dockerfile

@@ -1,24 +1,42 @@
-FROM python:3.11 AS builder
- 
+# Stage 1: Build
+FROM python:3.11-slim as builder
+
 WORKDIR /app
- 
-RUN python3 -m venv venv
-ENV VIRTUAL_ENV=/app/venv
-ENV PATH="$VIRTUAL_ENV/bin:$PATH"
- 
+
+# Copy requirements.txt
 COPY requirements.txt .
-RUN pip install -r requirements.txt
- 
-# Stage 2
-FROM python:3.11 AS runner
- 
+
+# Install dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy everything 
+COPY app app
+
+# Stage 2: Production
+FROM python:3.11-slim as production
+
 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" ]
+
+# Copy installed modules from the builder stage
+COPY --from=builder /usr/local/lib/python3.11/site-packages /usr/local/lib/python3.11/site-packages
+
+# Install playwright dependencies
+RUN python -m playwright install-deps && \
+    # Clean up
+    rm -f /usr/lib/x86_64-linux-gnu/libmfxhw* /usr/lib/x86_64-linux-gnu/mfx/* && \
+    # Create non-root user
+    useradd --home-dir /app --shell /bin/sh zaws && \
+    chown -R zaws:zaws .
+
+# Switch to non-root user
+USER zaws
+
+# Installfirefox inside non-root
+# Do not install firefox in root to avoid permission error otherwise you need to give zaws permission to access installation path
+RUN python -m playwright install firefox
+
+# Copy application code from the builder stage
+COPY --from=builder /app /app
+
+# Run application
+CMD ["python", "app/main.py"]

README.md

@@ -0,0 +1,11 @@
+---
+title: Zawshoot
+emoji: 📈
+colorFrom: pink
+colorTo: indigo
+sdk: docker
+pinned: false
+license: mit
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app/core/service/playwright/__init__.py

@@ -0,0 +1,5 @@
+# ruff: noqa: F401
+from .playwright_context import (
+    AsyncPlaywrightContext,
+    TimeoutError,
+)

app/core/service/playwright/models.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import BaseModel, Field, HttpUrl, ValidationInfo, field_validator
+
+
+class ViewPortModel(BaseModel):
+    width: int = 1280
+    height: int = 720
+
+
+class PageModel(BaseModel):
+    color_scheme: Literal["light", "dark", "no-preference"] | None = "no-preference"
+    java_script_enabled: bool | None = True
+    viewport: ViewPortModel | None = None
+    proxy: dict | None = None
+    no_viewport: bool | None = False
+
+
+class GetContentModel(BaseModel):
+    url: HttpUrl
+    query_selector: str | None = None
+    wait_selector: bool | None = False
+    ms_delay: float = Field(default=0.0, gt=-1, lt=15_000.1)
+
+    @field_validator("wait_selector")
+    @classmethod
+    def check_query_selector(cls, v: bool | None, info: ValidationInfo) -> bool | None:
+        if v is True and info.data.get("query_selector") is None:
+            msg = "wait_selector cannot be set to True without specifying a query_selector"
+            raise ValueError(
+                msg,
+            )
+        return v
+
+
+class ScreenshotModel(GetContentModel):
+    full_page: bool | None = False

app/core/service/playwright/playwright_context.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, AsyncContextManager
+
+from playwright.async_api import (  # noqa: F401
+    Browser,
+    BrowserContext,
+    Page,
+    TimeoutError,
+    async_playwright,
+)
+
+from .models import BrowserModel, GetContentModel, PageModel, ScreenshotModel  # noqa: TCH001
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+
+class AsyncPlaywrightContext:
+    async def __aenter__(self) -> AsyncContextManager:
+        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": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/83.0.4103.61 Safari/537.36",  # noqa: E501
+            },
+        )
+        return self
+
+    async def new_browser_page(self, browser: Browser, page_model: PageModel) -> Page:
+        return await browser.new_page(
+            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,
+        )
+
+    async def screenshot(
+        self,
+        screenshot_model: ScreenshotModel,
+        page_model: PageModel,
+    ) -> bytes:
+        page = await self.new_browser_page(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:
+            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,
+        browser_model: BrowserModel,
+    ) -> str:
+        browser_context = await self.new_browser_context(
+            browser=self.browser,
+            browser_model=browser_model,
+        )
+        page = browser_context.new_page()
+        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 wait_locator:
+            await wait_locator.wait_for()
+        html = page.content()
+        await page.close()
+        return html
+
+    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()

{screenshot → app}/main.py

@@ -1,10 +1,12 @@
+import os
+
 import uvicorn
 from fastapi import FastAPI
-from routers.screenshot import routers as screenshot_router
+from routers.screenshot import router 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
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", default="7860")))  # noqa: S104

app/routers/screenshot/__init__.py

@@ -0,0 +1,6 @@
+from fastapi import APIRouter
+
+from .screenshot import router as screenshot_router
+
+router = APIRouter()
+router.include_router(screenshot_router)

app/routers/screenshot/screenshot.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import logging
+
+from core.service.playwright import AsyncPlaywrightContext, TimeoutError
+from core.service.playwright.models import PageModel, ScreenshotModel
+from fastapi import APIRouter, HTTPException
+from fastapi.responses import Response
+
+router = APIRouter()
+playwright_context = AsyncPlaywrightContext()
+
+
+@router.post("/screenshot", tags=["V1"])
+async def screenshot(
+    screenshot_model: ScreenshotModel,
+    page_model: PageModel | None = None,
+) -> Response:
+    page_model = page_model if page_model else PageModel()
+    async with playwright_context as sc:
+        try:
+            response = await sc.screenshot(screenshot_model=screenshot_model, page_model=page_model)
+            return Response(content=response, media_type="image/png")
+        except TimeoutError as e:
+            status_code = 504 if isinstance(e, TimeoutError) else 400
+            detail = f"An error occurred while generating the screenshot: {e}"
+            raise HTTPException(status_code=status_code, detail=detail) from e
+        except Exception as e:
+            logging.exception("Unhandled error in screenshot generation")
+            raise HTTPException(status_code=500, detail="An unexpected error occurred") from e

pyproject.toml

@@ -5,7 +5,9 @@ line-length = 100
 select = ["ALL"]
 ignore = [
   "CPY001", # copyright above code
-  "D", # sphinx not support 
+  "ANN101", # self annotations
+  "ANN102", # cls annotations
+  "D",      # sphinx not support 
   ]
 
 [tool.mypy]
@@ -14,5 +16,5 @@ show_error_codes = true
 no_implicit_optional = true
 warn_return_any = true
 warn_unused_ignores = true
-exclude = ["tests"]
-python_version = "3.10"
+python_version = "3.10"
+plugins = ["pydantic.mypy"]

requirements-dev.txt

@@ -0,0 +1 @@
+ruff==0.3.3

requirements.txt

@@ -0,0 +1,4 @@
+pydantic==2.6.4
+fastapi==0.110.0
+playwright==1.42.0
+uvicorn==0.28.0

screenshot/routers/screenshot.py

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