Upload dom_element_extractor.py

utils/dom_element_extractor.py

@@ -0,0 +1,513 @@
+"""
+DOM Element Extractor
+Extracts UI elements from websites using Playwright.
+Gets computed styles and element properties for comparison with Figma.
+"""
+
+import asyncio
+from typing import Dict, List, Optional, Any, Tuple
+from .element_schema import (
+    UIElement, ElementType, ElementBounds, ElementStyles,
+    DOM_TYPE_MAP
+)
+
+
+# JavaScript to extract elements and their computed styles
+EXTRACTION_SCRIPT = """
+() => {
+    const results = [];
+    
+    // Selectors for elements we care about (checkout-focused)
+    const selectors = [
+        'button',
+        'input',
+        'textarea',
+        'select',
+        'a',
+        'h1', 'h2', 'h3', 'h4', 'h5', 'h6',
+        'p',
+        'label',
+        'span',
+        'img',
+        'svg',
+        '[role="button"]',
+        '[role="link"]',
+        '[role="textbox"]',
+        '[role="checkbox"]',
+        '[role="radio"]',
+        // Common checkout classes
+        '.btn', '.button',
+        '.input', '.field',
+        '.price', '.total', '.amount',
+        '.card', '.form-group',
+        // Text classes (Tailwind)
+        '.text-muted-foreground',
+        '.text-gray-500', '.text-gray-600', '.text-gray-700',
+        // Data attributes
+        '[data-testid]',
+        '[data-cy]'
+    ];
+    
+    // Get all matching elements
+    const elements = new Set();
+    selectors.forEach(selector => {
+        try {
+            document.querySelectorAll(selector).forEach(el => elements.add(el));
+        } catch(e) {}
+    });
+    
+    // Also get important containers (forms, main sections)
+    document.querySelectorAll('form, main, [role="main"], .checkout, .cart, .payment').forEach(el => elements.add(el));
+    
+    // Process each element
+    let index = 0;
+    elements.forEach(el => {
+        // Skip hidden elements
+        const style = window.getComputedStyle(el);
+        if (style.display === 'none' || style.visibility === 'hidden' || style.opacity === '0') {
+            return;
+        }
+        
+        // Get bounding rect
+        const rect = el.getBoundingClientRect();
+        
+        // Skip elements outside viewport or too small
+        if (rect.width < 5 || rect.height < 5) {
+            return;
+        }
+        
+        // Skip elements way off screen
+        if (rect.top > 10000 || rect.left > 5000) {
+            return;
+        }
+        
+        // Get element info
+        const tagName = el.tagName.toLowerCase();
+        const id = el.id || '';
+        const className = el.className || '';
+        const classStr = typeof className === 'string' ? className : '';
+        
+        // Determine name (prefer aria-label, then id, then class, then tag)
+        let name = el.getAttribute('aria-label') || 
+                   el.getAttribute('data-testid') ||
+                   el.getAttribute('placeholder') ||
+                   id ||
+                   classStr.split(' ')[0] ||
+                   tagName;
+        
+        // Get text content (direct text only, not nested)
+        let textContent = '';
+        if (['BUTTON', 'A', 'LABEL', 'H1', 'H2', 'H3', 'H4', 'H5', 'H6', 'P', 'SPAN'].includes(el.tagName)) {
+            // Get direct text content
+            textContent = Array.from(el.childNodes)
+                .filter(node => node.nodeType === Node.TEXT_NODE)
+                .map(node => node.textContent.trim())
+                .join(' ')
+                .trim();
+            
+            // If no direct text, get innerText but truncate
+            if (!textContent && el.innerText) {
+                textContent = el.innerText.substring(0, 100).trim();
+            }
+        }
+        
+        // Use text content as name if element has meaningful text
+        if (textContent && textContent.length > 2 && textContent.length < 80) {
+            name = textContent;
+        }
+        
+        // Get placeholder for inputs
+        const placeholder = el.getAttribute('placeholder') || '';
+        
+        // Get input type
+        const inputType = el.getAttribute('type') || '';
+        
+        // IMPORTANT: Get the ACTUAL computed color, not inherited
+        // Force a reflow to ensure styles are computed
+        el.offsetHeight;
+        
+        // Re-get computed style after reflow
+        const finalStyle = window.getComputedStyle(el);
+        
+        // Extract computed styles
+        const computedStyles = {
+            backgroundColor: finalStyle.backgroundColor,
+            color: finalStyle.color,
+            borderColor: finalStyle.borderColor,
+            borderWidth: finalStyle.borderWidth,
+            borderRadius: finalStyle.borderRadius,
+            borderStyle: finalStyle.borderStyle,
+            fontFamily: finalStyle.fontFamily,
+            fontSize: finalStyle.fontSize,
+            fontWeight: finalStyle.fontWeight,
+            lineHeight: finalStyle.lineHeight,
+            textAlign: finalStyle.textAlign,
+            letterSpacing: finalStyle.letterSpacing,
+            paddingTop: finalStyle.paddingTop,
+            paddingRight: finalStyle.paddingRight,
+            paddingBottom: finalStyle.paddingBottom,
+            paddingLeft: finalStyle.paddingLeft,
+            opacity: finalStyle.opacity,
+            boxShadow: finalStyle.boxShadow
+        };
+        
+        // Debug: Log color for text elements
+        if (textContent && textContent.includes('Complete your purchase')) {
+            console.log('Found element:', textContent, 'Color:', finalStyle.color);
+        }
+        
+        results.push({
+            index: index++,
+            tagName: tagName,
+            id: id,
+            className: classStr,
+            name: name,
+            textContent: textContent,
+            placeholder: placeholder,
+            inputType: inputType,
+            bounds: {
+                x: rect.left + window.scrollX,
+                y: rect.top + window.scrollY,
+                width: rect.width,
+                height: rect.height
+            },
+            styles: computedStyles,
+            isInteractive: ['BUTTON', 'A', 'INPUT', 'TEXTAREA', 'SELECT'].includes(el.tagName) ||
+                          el.getAttribute('role') === 'button' ||
+                          el.onclick !== null
+        });
+    });
+    
+    return results;
+}
+"""
+
+
+class DOMElementExtractor:
+    """
+    Extracts UI elements from a webpage using Playwright.
+    Gets computed styles that match what's actually rendered.
+    """
+    
+    def __init__(self):
+        self.elements: List[UIElement] = []
+    
+    async def extract_from_page_async(
+        self,
+        page,  # Playwright page object
+        viewport: str = "desktop"
+    ) -> List[UIElement]:
+        """
+        Extract elements from a Playwright page.
+        
+        Args:
+            page: Playwright page object (already navigated to URL)
+            viewport: "desktop" or "mobile"
+            
+        Returns:
+            List of UIElement objects
+        """
+        self.elements = []
+        
+        # Execute extraction script
+        raw_elements = await page.evaluate(EXTRACTION_SCRIPT)
+        
+        # Convert to UIElement objects
+        for raw in raw_elements:
+            element = self._create_element(raw, viewport)
+            if element:
+                self.elements.append(element)
+        
+        print(f"     📊 Extracted {len(self.elements)} elements from {viewport} page")
+        return self.elements
+    
+    def _create_element(self, raw: Dict, viewport: str) -> Optional[UIElement]:
+        """
+        Create a UIElement from raw extracted data.
+        """
+        # Create bounds
+        bounds_data = raw.get("bounds", {})
+        bounds = ElementBounds(
+            x=bounds_data.get("x", 0),
+            y=bounds_data.get("y", 0),
+            width=bounds_data.get("width", 0),
+            height=bounds_data.get("height", 0)
+        )
+        
+        # Determine element type
+        element_type = self._determine_element_type(raw)
+        
+        # Extract styles
+        styles = self._extract_styles(raw.get("styles", {}))
+        
+        # Create element
+        return UIElement(
+            id=f"dom_{raw.get('index', 0)}_{raw.get('tagName', 'unknown')}",
+            element_type=element_type,
+            name=raw.get("name", ""),
+            bounds=bounds,
+            styles=styles,
+            text_content=raw.get("textContent", "") or None,
+            placeholder=raw.get("placeholder", "") or None,
+            source="website",
+            original_type=raw.get("tagName", ""),
+            is_interactive=raw.get("isInteractive", False),
+            input_type=raw.get("inputType", "") or None
+        )
+    
+    def _determine_element_type(self, raw: Dict) -> ElementType:
+        """
+        Determine semantic element type from DOM element.
+        """
+        tag = raw.get("tagName", "").lower()
+        class_name = raw.get("className", "").lower()
+        text = raw.get("textContent", "").lower()
+        input_type = raw.get("inputType", "").lower()
+        
+        # Check tag mapping first
+        if tag in DOM_TYPE_MAP:
+            base_type = DOM_TYPE_MAP[tag]
+            
+            # Refine input types
+            if tag == "input":
+                if input_type == "checkbox":
+                    return ElementType.CHECKBOX
+                elif input_type == "radio":
+                    return ElementType.RADIO
+                elif input_type == "submit" or input_type == "button":
+                    return ElementType.BUTTON
+                return ElementType.INPUT
+            
+            return base_type
+        
+        # Check class names for hints
+        if any(kw in class_name for kw in ["btn", "button"]):
+            return ElementType.BUTTON
+        if any(kw in class_name for kw in ["input", "field", "textbox"]):
+            return ElementType.INPUT
+        if any(kw in class_name for kw in ["price", "amount", "total", "cost"]):
+            return ElementType.PRICE
+        if any(kw in class_name for kw in ["card"]):
+            return ElementType.CARD
+        if any(kw in class_name for kw in ["icon", "ico"]):
+            return ElementType.ICON
+        if any(kw in class_name for kw in ["badge", "tag", "chip"]):
+            return ElementType.BADGE
+        
+        # Check text content for price detection
+        if "$" in text or "€" in text or "£" in text:
+            return ElementType.PRICE
+        
+        return ElementType.UNKNOWN
+    
+    def _extract_styles(self, raw_styles: Dict) -> ElementStyles:
+        """
+        Convert computed CSS styles to ElementStyles.
+        """
+        styles = ElementStyles()
+        
+        # Background color
+        bg = raw_styles.get("backgroundColor", "")
+        if bg and bg != "rgba(0, 0, 0, 0)" and bg != "transparent":
+            styles.background_color = self._css_color_to_hex(bg)
+        
+        # Text color
+        color = raw_styles.get("color", "")
+        if color:
+            styles.text_color = self._css_color_to_hex(color)
+        
+        # Border color
+        border_color = raw_styles.get("borderColor", "")
+        if border_color:
+            styles.border_color = self._css_color_to_hex(border_color)
+        
+        # Border width (parse "1px" -> 1.0)
+        border_width = raw_styles.get("borderWidth", "")
+        if border_width:
+            styles.border_width = self._parse_px_value(border_width)
+        
+        # Border radius
+        border_radius = raw_styles.get("borderRadius", "")
+        if border_radius:
+            styles.border_radius = self._parse_px_value(border_radius)
+        
+        # Border style
+        border_style = raw_styles.get("borderStyle", "")
+        if border_style and border_style != "none":
+            styles.border_style = border_style
+        
+        # Font family (clean up quotes)
+        font_family = raw_styles.get("fontFamily", "")
+        if font_family:
+            # Take first font in the stack
+            styles.font_family = font_family.split(",")[0].strip().strip('"\'')
+        
+        # Font size
+        font_size = raw_styles.get("fontSize", "")
+        if font_size:
+            styles.font_size = self._parse_px_value(font_size)
+        
+        # Font weight
+        font_weight = raw_styles.get("fontWeight", "")
+        if font_weight:
+            try:
+                styles.font_weight = int(font_weight)
+            except ValueError:
+                # Handle named weights
+                weight_map = {"normal": 400, "bold": 700, "lighter": 300, "bolder": 700}
+                styles.font_weight = weight_map.get(font_weight, 400)
+        
+        # Line height
+        line_height = raw_styles.get("lineHeight", "")
+        if line_height and line_height != "normal":
+            styles.line_height = self._parse_px_value(line_height)
+        
+        # Text align
+        text_align = raw_styles.get("textAlign", "")
+        if text_align:
+            styles.text_align = text_align
+        
+        # Letter spacing
+        letter_spacing = raw_styles.get("letterSpacing", "")
+        if letter_spacing and letter_spacing != "normal":
+            styles.letter_spacing = self._parse_px_value(letter_spacing)
+        
+        # Padding
+        styles.padding_top = self._parse_px_value(raw_styles.get("paddingTop", ""))
+        styles.padding_right = self._parse_px_value(raw_styles.get("paddingRight", ""))
+        styles.padding_bottom = self._parse_px_value(raw_styles.get("paddingBottom", ""))
+        styles.padding_left = self._parse_px_value(raw_styles.get("paddingLeft", ""))
+        
+        # Opacity
+        opacity = raw_styles.get("opacity", "")
+        if opacity:
+            try:
+                styles.opacity = float(opacity)
+            except ValueError:
+                pass
+        
+        # Box shadow
+        box_shadow = raw_styles.get("boxShadow", "")
+        if box_shadow and box_shadow != "none":
+            styles.box_shadow = box_shadow
+        
+        return styles
+    
+    def _css_color_to_hex(self, css_color: str) -> Optional[str]:
+        """
+        Convert CSS color (rgb, rgba, hex) to hex string.
+        """
+        if not css_color:
+            return None
+        
+        css_color = css_color.strip()
+        
+        # Already hex
+        if css_color.startswith("#"):
+            return css_color.upper()
+        
+        # rgb() or rgba()
+        if css_color.startswith("rgb"):
+            try:
+                # Extract numbers
+                import re
+                numbers = re.findall(r'[\d.]+', css_color)
+                if len(numbers) >= 3:
+                    r, g, b = int(float(numbers[0])), int(float(numbers[1])), int(float(numbers[2]))
+                    hex_color = f"#{r:02X}{g:02X}{b:02X}"
+                    return hex_color
+            except Exception as e:
+                print(f"        ⚠️ Color parse error: {css_color} - {e}")
+        
+        return None
+    
+    def _parse_px_value(self, value: str) -> Optional[float]:
+        """
+        Parse CSS pixel value (e.g., "16px") to float.
+        """
+        if not value:
+            return None
+        
+        try:
+            # Remove 'px' and convert
+            return float(value.replace("px", "").strip())
+        except ValueError:
+            return None
+    
+    def get_interactive_elements(self) -> List[UIElement]:
+        """Get only interactive elements."""
+        return [e for e in self.elements if e.is_interactive]
+    
+    def get_elements_by_type(self, element_type: ElementType) -> List[UIElement]:
+        """Get elements of a specific type."""
+        return [e for e in self.elements if e.element_type == element_type]
+    
+    def summarize(self) -> Dict[str, Any]:
+        """Get a summary of extracted elements."""
+        type_counts = {}
+        for element in self.elements:
+            type_name = element.element_type.value
+            type_counts[type_name] = type_counts.get(type_name, 0) + 1
+        
+        return {
+            "total_elements": len(self.elements),
+            "interactive_elements": len(self.get_interactive_elements()),
+            "by_type": type_counts
+        }
+
+
+async def extract_dom_elements_async(
+    url: str,
+    viewport_width: int = 1440,
+    viewport_height: int = 900,
+    viewport_name: str = "desktop"
+) -> Tuple[List[UIElement], Dict]:
+    """
+    Extract elements from a URL.
+    
+    Args:
+        url: Website URL to extract from
+        viewport_width: Viewport width in pixels
+        viewport_height: Viewport height in pixels
+        viewport_name: "desktop" or "mobile"
+        
+    Returns:
+        Tuple of (list of UIElements, summary dict)
+    """
+    from playwright.async_api import async_playwright
+    
+    async with async_playwright() as p:
+        browser = await p.chromium.launch(headless=True)
+        
+        try:
+            page = await browser.new_page(viewport={"width": viewport_width, "height": viewport_height})
+            
+            # Set mobile user agent if needed
+            if viewport_name == "mobile":
+                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(url, wait_until="networkidle", timeout=60000)
+            await page.wait_for_timeout(2000)  # Wait for any animations
+            
+            extractor = DOMElementExtractor()
+            elements = await extractor.extract_from_page_async(page, viewport_name)
+            summary = extractor.summarize()
+            
+            return elements, summary
+            
+        finally:
+            await browser.close()
+
+
+def extract_dom_elements(
+    url: str,
+    viewport_width: int = 1440,
+    viewport_height: int = 900,
+    viewport_name: str = "desktop"
+) -> Tuple[List[UIElement], Dict]:
+    """
+    Synchronous wrapper for DOM element extraction.
+    """
+    return asyncio.run(extract_dom_elements_async(url, viewport_width, viewport_height, viewport_name))