File size: 4,266 Bytes
cfec14d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
"""
Website Screenshot Capturer
Uses Playwright to capture full-page screenshots of websites.
"""

import asyncio
from pathlib import Path
from typing import Dict, Tuple


async def _capture_screenshots_async(
    website_url: str,
    output_dir: str,
    execution_id: str,
    desktop_width: int = 1440,
    mobile_width: int = 375
) -> Tuple[Dict[str, str], Dict[str, Dict[str, int]]]:
    """
    Async function to capture website screenshots.
    
    Captures full-page screenshots at desktop and mobile viewports.
    """
    from playwright.async_api import async_playwright
    
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    screenshots = {}
    dimensions = {}
    
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=True)
        
        try:
            # Desktop capture
            print(f"  📱 Capturing desktop ({desktop_width}px width, full height)...")
            page = await browser.new_page(viewport={"width": desktop_width, "height": 1080})
            await page.goto(website_url, wait_until="networkidle", timeout=60000)
            
            # Wait for any lazy-loaded content
            await page.wait_for_timeout(2000)
            
            # Get full page height
            desktop_height = await page.evaluate("() => document.documentElement.scrollHeight")
            print(f"     ℹ️  Desktop full height: {desktop_height}px")
            
            # Set viewport to full height and capture
            await page.set_viewport_size({"width": desktop_width, "height": desktop_height})
            desktop_path = f"{output_dir}/desktop_{execution_id}.png"
            await page.screenshot(path=desktop_path, full_page=True)
            
            screenshots["desktop"] = desktop_path
            dimensions["desktop"] = {"width": desktop_width, "height": desktop_height}
            print(f"     ✓ Saved: {desktop_path}")
            
            await page.close()
            
            # Mobile capture
            print(f"  📱 Capturing mobile ({mobile_width}px width, full height)...")
            page = await browser.new_page(viewport={"width": mobile_width, "height": 812})
            
            # Set mobile user agent
            await page.set_extra_http_headers({
                "User-Agent": "Mozilla/5.0 (iPhone; CPU iPhone OS 15_0 like Mac OS X) AppleWebKit/605.1.15"
            })
            
            await page.goto(website_url, wait_until="networkidle", timeout=60000)
            await page.wait_for_timeout(2000)
            
            # Get full page height
            mobile_height = await page.evaluate("() => document.documentElement.scrollHeight")
            print(f"     ℹ️  Mobile full height: {mobile_height}px")
            
            # Set viewport to full height and capture
            await page.set_viewport_size({"width": mobile_width, "height": mobile_height})
            mobile_path = f"{output_dir}/mobile_{execution_id}.png"
            await page.screenshot(path=mobile_path, full_page=True)
            
            screenshots["mobile"] = mobile_path
            dimensions["mobile"] = {"width": mobile_width, "height": mobile_height}
            print(f"     ✓ Saved: {mobile_path}")
            
            await page.close()
            
        finally:
            await browser.close()
    
    return screenshots, dimensions


def capture_website_screenshots(
    website_url: str,
    output_dir: str,
    execution_id: str,
    desktop_width: int = 1440,
    mobile_width: int = 375
) -> Tuple[Dict[str, str], Dict[str, Dict[str, int]]]:
    """
    Synchronous wrapper for capturing website screenshots.
    
    Args:
        website_url: URL of the website to capture
        output_dir: Directory to save screenshots
        execution_id: Unique ID for this execution
        desktop_width: Desktop viewport width (default 1440)
        mobile_width: Mobile viewport width (default 375)
    
    Returns:
        Tuple of (screenshot_paths_dict, dimensions_dict)
    """
    return asyncio.run(
        _capture_screenshots_async(
            website_url,
            output_dir,
            execution_id,
            desktop_width,
            mobile_width
        )
    )