fetch_screen_node.py

"""
fetch_screen_node module
"""

from typing import List, Optional

from playwright.sync_api import sync_playwright

from .base_node import BaseNode


class FetchScreenNode(BaseNode):
    """
    FetchScreenNode captures screenshots from a given URL and stores the image data as bytes.
    """

    def __init__(
        self,
        input: str,
        output: List[str],
        node_config: Optional[dict] = None,
        node_name: str = "FetchScreen",
    ):
        super().__init__(node_name, "node", input, output, 2, node_config)
        self.url = node_config.get("link")

    def execute(self, state: dict) -> dict:
        """
        Captures screenshots from the input URL and stores them in the state dictionary as bytes.
        """
        self.logger.info(f"--- Executing {self.node_name} Node ---")

        with sync_playwright() as p:
            browser = p.chromium.launch()
            page = browser.new_page()
            page.goto(self.url)

            viewport_height = page.viewport_size["height"]

            screenshot_counter = 1

            screenshot_data_list = []

            def capture_screenshot(scroll_position, counter):
                page.evaluate(f"window.scrollTo(0, {scroll_position});")
                screenshot_data = page.screenshot()
                screenshot_data_list.append(screenshot_data)

            capture_screenshot(0, screenshot_counter)
            screenshot_counter += 1
            capture_screenshot(viewport_height, screenshot_counter)

            browser.close()

        state["link"] = self.url
        state["screenshots"] = screenshot_data_list

        return state