diff --git a/tools/itb/bin/itb_screenshot b/tools/itb/bin/itb_screenshot index 5440e10..160d957 100755 --- a/tools/itb/bin/itb_screenshot +++ b/tools/itb/bin/itb_screenshot @@ -2,196 +2,140 @@ import sys from selenium import webdriver from selenium.webdriver.common.by import By -from selenium.webdriver.remote.webelement import WebElement import xml.etree.ElementTree as ET -from typing import Dict, List, Set, Optional, Union -from dataclasses import dataclass - -@dataclass -class DOMNode: - element: WebElement - children: List['DOMNode'] - text: str - tag: str - location: Dict[str, int] - size: Dict[str, int] - interactions: List[str] - - def get_content(self) -> List[Union[ET.Element, str]]: - content = [] - - # Add text as plain string if present - if self.text: - content.append(self.text) - - # Add interactive element if present - if self.interactions: - attrs = { - "x": str(self.location["x"]), - "y": str(self.location["y"]), - "width": str(self.size["width"]), - "height": str(self.size["height"]) - } - - # Split interactions and placeholders - base_interactions = [] - placeholder = None - - for interaction in self.interactions: - if interaction.startswith("placeholder:"): - placeholder = interaction.split(":", 1)[1] - else: - base_interactions.append(interaction) - - attrs["interactions"] = ",".join(base_interactions) - if placeholder: - attrs["placeholder"] = placeholder - - interactive = ET.Element("interactive", attrs) - if self.text: - interactive.text = self.text - content.append(interactive) - - # Add all children's content - for child in self.children: - content.extend(child.get_content()) - - return content - -class DOMTree: - def __init__(self, root_element: WebElement): - print(f"\nBuilding DOM tree for {root_element.tag_name}") - self.root = self._build_tree(root_element) - - def _get_element_text(self, element: WebElement) -> str: - script = """ - let element = arguments[0]; - let text = ''; - for (let node of element.childNodes) { - if (node.nodeType === Node.TEXT_NODE) { - text += node.textContent; - } - } - return text.trim(); - """ - text = element.parent.execute_script(script, element) - if text: - print(f"Got text: {text[:50]}...") - return text - - def _get_interactions(self, element: WebElement) -> List[str]: - interactions = [] - - if (element.tag_name in {"button", "a"} or - element.get_attribute("role") in {"button", "link"} or - element.get_attribute("onclick") or - element.value_of_css_property("cursor") == "pointer"): - interactions.append("click") - - if (element.tag_name in {"input", "textarea"} or - element.get_attribute("contenteditable") == "true" or - element.get_attribute("role") == "textbox"): - interactions.append("text_input") - # Get placeholder if available - placeholder = element.get_attribute("placeholder") - if placeholder: - interactions.append(f"placeholder:{placeholder}") - - if (element.tag_name == "select" or - element.get_attribute("role") == "listbox" or - element.get_attribute("type") in {"checkbox", "radio"}): - interactions.append("select") - - script = """ - let el = arguments[0]; - return el.scrollHeight > el.clientHeight || - el.scrollWidth > el.clientWidth; - """ - if element.parent.execute_script(script, element): - interactions.append("scroll") - - if interactions: - print(f"Found interactions for {element.tag_name}: {interactions}") - return interactions - - def _build_tree(self, element: WebElement) -> DOMNode: - print(f"\nProcessing <{element.tag_name}>") - - if not element.is_displayed(): - print(f"Skipping invisible {element.tag_name}") - return None - - children = [] - for child in element.find_elements(By.XPATH, "./*"): - if child_node := self._build_tree(child): - children.append(child_node) - - text = self._get_element_text(element) - interactions = self._get_interactions(element) - - return DOMNode( - element=element, - children=children, - text=text, - tag=element.tag_name, - location=element.location, - size=element.size, - interactions=interactions - ) class ScreenshotGenerator: def __init__(self, debug_port: int): - print(f"\nInitializing ScreenshotGenerator with port {debug_port}") options = webdriver.ChromeOptions() options.add_experimental_option("debuggerAddress", f"127.0.0.1:{debug_port}") self.driver = webdriver.Chrome(options=options) - def get_viewport_metrics(self) -> Dict[str, str]: - metrics = { - "scroll_x": self.driver.execute_script("return window.pageXOffset;"), - "scroll_y": self.driver.execute_script("return window.pageYOffset;"), - "width": self.driver.execute_script("return document.documentElement.clientWidth;"), - "height": self.driver.execute_script("return document.documentElement.clientHeight;"), - "max_scroll_y": self.driver.execute_script( - "return Math.max(document.documentElement.scrollHeight - window.innerHeight, 0);" - ), - "min_scroll_x": self.driver.execute_script( - "return Math.max(document.documentElement.scrollLeft, 0);" - ) + def get_page_content(self) -> dict: + script = """ + function processElement(element) { + if (!element.offsetParent && element !== document.body) return null; + + const rect = element.getBoundingClientRect(); + const styles = window.getComputedStyle(element); + + const text = Array.from(element.childNodes) + .filter(node => node.nodeType === Node.TEXT_NODE) + .map(node => node.textContent.trim()) + .join(' '); + + const interactions = []; + if (element.tagName === 'BUTTON' || element.tagName === 'A' || + element.getAttribute('role') === 'button' || + element.getAttribute('onclick') || + styles.cursor === 'pointer') { + interactions.push('click'); + } + + if (element.tagName === 'INPUT' || element.tagName === 'TEXTAREA' || + element.getAttribute('contenteditable') === 'true') { + interactions.push('text_input'); + const placeholder = element.getAttribute('placeholder'); + if (placeholder) interactions.push('placeholder:' + placeholder); + } + + if (element.tagName === 'SELECT' || + element.getAttribute('role') === 'listbox' || + element.getAttribute('type') === 'checkbox' || + element.getAttribute('type') === 'radio') { + interactions.push('select'); + } + + if (element.scrollHeight > element.clientHeight || + element.scrollWidth > element.clientWidth) { + interactions.push('scroll'); + } + + const children = Array.from(element.children) + .map(child => processElement(child)) + .filter(child => child !== null); + + return { + tag: element.tagName.toLowerCase(), + text, + location: {x: Math.round(rect.left), y: Math.round(rect.top)}, + size: {width: Math.round(rect.width), height: Math.round(rect.height)}, + interactions, + children + }; } - print(f"\nViewport metrics: {metrics}") - return {k: str(v) for k, v in metrics.items()} + + const body = document.body; + return { + viewport: { + scroll_x: window.pageXOffset, + scroll_y: window.pageYOffset, + width: document.documentElement.clientWidth, + height: document.documentElement.clientHeight, + max_scroll_y: Math.max(document.documentElement.scrollHeight - window.innerHeight, 0), + min_scroll_x: Math.max(document.documentElement.scrollLeft, 0) + }, + content: body ? processElement(body) : null + }; + """ + return self.driver.execute_script(script) + + def format_xml(self, xml_str: str) -> str: + root = ET.fromstring(xml_str) + ET.indent(ET.ElementTree(root), space=" ", level=0) + return ET.tostring(root, encoding='unicode', method='xml') def generate_xml(self) -> str: - metrics = self.get_viewport_metrics() - metrics_str = " ".join(f'{k}="{v}"' for k, v in metrics.items()) + page_data = self.get_page_content() - output = [f''] - seen_texts = set() - last_was_text = False + viewport = ET.Element("viewport") + for k, v in page_data['viewport'].items(): + viewport.set(k, str(v)) - body = self.driver.find_element(By.TAG_NAME, "body") - dom_tree = DOMTree(body) - - for content in dom_tree.root.get_content(): - if isinstance(content, str): - if content not in seen_texts: - if last_was_text: - output.append("\n") - output.append(content) - seen_texts.add(content) - last_was_text = True - else: - output.append(ET.tostring(content, encoding="unicode")) - last_was_text = False + def process_node(node, parent): + if node is None: + return - output.append('') - return "".join(output) + if node['text']: + if len(parent) == 0: + parent.text = (parent.text or '') + node['text'] + '\n' + else: + parent[-1].tail = (parent[-1].tail or '') + node['text'] + '\n' + + if node['interactions']: + interactive = ET.SubElement(parent, "interactive") + interactive.set("x", str(node['location']['x'])) + interactive.set("y", str(node['location']['y'])) + interactive.set("width", str(node['size']['width'])) + interactive.set("height", str(node['size']['height'])) + interactive.set("interactions", ",".join( + i for i in node['interactions'] if not i.startswith('placeholder:') + )) + + placeholder = next( + (i.split(':', 1)[1] for i in node['interactions'] + if i.startswith('placeholder:')), + None + ) + if placeholder: + interactive.set('placeholder', placeholder) + + if node['text']: + interactive.text = node['text'] + + for child in node['children']: + process_node(child, parent) + + process_node(page_data['content'], viewport) + return self.format_xml(ET.tostring(viewport, encoding='unicode')) def main(debug_port: int): - generator = ScreenshotGenerator(debug_port) - print("\nFinal XML output:") - print(generator.generate_xml()) + try: + generator = ScreenshotGenerator(debug_port) + xml_output = generator.generate_xml() + print(xml_output) + except Exception as e: + print(f"Error: {e}", file=sys.stderr) + sys.exit(1) if __name__ == "__main__": if len(sys.argv) != 2: