First version of itb_screenshot
This commit is contained in:
@@ -2,196 +2,140 @@
|
|||||||
import sys
|
import sys
|
||||||
from selenium import webdriver
|
from selenium import webdriver
|
||||||
from selenium.webdriver.common.by import By
|
from selenium.webdriver.common.by import By
|
||||||
from selenium.webdriver.remote.webelement import WebElement
|
|
||||||
import xml.etree.ElementTree as ET
|
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:
|
class ScreenshotGenerator:
|
||||||
def __init__(self, debug_port: int):
|
def __init__(self, debug_port: int):
|
||||||
print(f"\nInitializing ScreenshotGenerator with port {debug_port}")
|
|
||||||
options = webdriver.ChromeOptions()
|
options = webdriver.ChromeOptions()
|
||||||
options.add_experimental_option("debuggerAddress", f"127.0.0.1:{debug_port}")
|
options.add_experimental_option("debuggerAddress", f"127.0.0.1:{debug_port}")
|
||||||
self.driver = webdriver.Chrome(options=options)
|
self.driver = webdriver.Chrome(options=options)
|
||||||
|
|
||||||
def get_viewport_metrics(self) -> Dict[str, str]:
|
def get_page_content(self) -> dict:
|
||||||
metrics = {
|
script = """
|
||||||
"scroll_x": self.driver.execute_script("return window.pageXOffset;"),
|
function processElement(element) {
|
||||||
"scroll_y": self.driver.execute_script("return window.pageYOffset;"),
|
if (!element.offsetParent && element !== document.body) return null;
|
||||||
"width": self.driver.execute_script("return document.documentElement.clientWidth;"),
|
|
||||||
"height": self.driver.execute_script("return document.documentElement.clientHeight;"),
|
const rect = element.getBoundingClientRect();
|
||||||
"max_scroll_y": self.driver.execute_script(
|
const styles = window.getComputedStyle(element);
|
||||||
"return Math.max(document.documentElement.scrollHeight - window.innerHeight, 0);"
|
|
||||||
),
|
const text = Array.from(element.childNodes)
|
||||||
"min_scroll_x": self.driver.execute_script(
|
.filter(node => node.nodeType === Node.TEXT_NODE)
|
||||||
"return Math.max(document.documentElement.scrollLeft, 0);"
|
.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:
|
def generate_xml(self) -> str:
|
||||||
metrics = self.get_viewport_metrics()
|
page_data = self.get_page_content()
|
||||||
metrics_str = " ".join(f'{k}="{v}"' for k, v in metrics.items())
|
|
||||||
|
|
||||||
output = [f'<viewport {metrics_str}>']
|
viewport = ET.Element("viewport")
|
||||||
seen_texts = set()
|
for k, v in page_data['viewport'].items():
|
||||||
last_was_text = False
|
viewport.set(k, str(v))
|
||||||
|
|
||||||
body = self.driver.find_element(By.TAG_NAME, "body")
|
def process_node(node, parent):
|
||||||
dom_tree = DOMTree(body)
|
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
|
|
||||||
else:
|
|
||||||
output.append(ET.tostring(content, encoding="unicode"))
|
|
||||||
last_was_text = False
|
|
||||||
|
|
||||||
output.append('</viewport>')
|
if node['text']:
|
||||||
return "".join(output)
|
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):
|
def main(debug_port: int):
|
||||||
generator = ScreenshotGenerator(debug_port)
|
try:
|
||||||
print("\nFinal XML output:")
|
generator = ScreenshotGenerator(debug_port)
|
||||||
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 __name__ == "__main__":
|
||||||
if len(sys.argv) != 2:
|
if len(sys.argv) != 2:
|
||||||
|
|||||||
Reference in New Issue
Block a user