First version of itb_screenshot
This commit is contained in:
@@ -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');
|
||||
}
|
||||
print(f"\nViewport metrics: {metrics}")
|
||||
return {k: str(v) for k, v in metrics.items()}
|
||||
|
||||
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
|
||||
};
|
||||
}
|
||||
|
||||
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'<viewport {metrics_str}>']
|
||||
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)
|
||||
def process_node(node, parent):
|
||||
if node is None:
|
||||
return
|
||||
|
||||
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
|
||||
if node['text']:
|
||||
if len(parent) == 0:
|
||||
parent.text = (parent.text or '') + node['text'] + '\n'
|
||||
else:
|
||||
output.append(ET.tostring(content, encoding="unicode"))
|
||||
last_was_text = False
|
||||
parent[-1].tail = (parent[-1].tail or '') + node['text'] + '\n'
|
||||
|
||||
output.append('</viewport>')
|
||||
return "".join(output)
|
||||
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):
|
||||
try:
|
||||
generator = ScreenshotGenerator(debug_port)
|
||||
print("\nFinal XML output:")
|
||||
print(generator.generate_xml())
|
||||
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:
|
||||
|
||||
Reference in New Issue
Block a user