from __future__ import annotations import re import xml.etree.ElementTree as ET from typing import Any from core.models import Bounds, SceneElement _ANDROID_BOUNDS = re.compile(r"\[(?P-?\d+),(?P-?\d+)\]\[(?P-?\d+),(?P-?\d+)\]") def parse_ui_tree(raw_tree: Any) -> list[SceneElement]: if raw_tree is None: return [] if isinstance(raw_tree, str): return _parse_xml(raw_tree) if isinstance(raw_tree, dict): elements: list[SceneElement] = [] _parse_dict_node(raw_tree, elements) return _with_stable_ids(elements, "ui") if isinstance(raw_tree, list): elements = [] for node in raw_tree: if isinstance(node, dict): _parse_dict_node(node, elements) return _with_stable_ids(elements, "ui") return [] def _parse_xml(raw_tree: str) -> list[SceneElement]: if not raw_tree.strip(): return [] root = ET.fromstring(raw_tree) elements = [] for node in root.iter(): bounds = _bounds_from_attrs(node.attrib) if bounds is None or bounds.width <= 0 or bounds.height <= 0: continue elements.append( SceneElement( id="", type=_normalize_type(_local_name(node.tag), node.attrib), text=_text_from_attrs(node.attrib), bounds=bounds, confidence=1.0, source="ui", ) ) return _with_stable_ids(elements, "ui") def _parse_dict_node(node: dict[str, Any], elements: list[SceneElement]) -> None: attrs = dict(node) children = attrs.pop("children", None) or attrs.pop("nodes", None) or [] bounds = _bounds_from_attrs(attrs) if bounds and bounds.width > 0 and bounds.height > 0: elements.append( SceneElement( id="", type=_normalize_type( str(attrs.get("type") or attrs.get("class") or "unknown"), attrs, ), text=_text_from_attrs(attrs), bounds=bounds, confidence=1.0, source="ui", ) ) for child in children: if isinstance(child, dict): _parse_dict_node(child, elements) def _with_stable_ids(elements: list[SceneElement], prefix: str) -> list[SceneElement]: return [ SceneElement( id=element.id or f"{prefix}-{index:03d}", type=element.type, text=element.text, bounds=element.bounds, confidence=element.confidence, source=element.source, ) for index, element in enumerate(elements) ] def _bounds_from_attrs(attrs: dict[str, Any]) -> Bounds | None: if all(key in attrs for key in ("x", "y", "width", "height")): return Bounds( float(attrs["x"]), float(attrs["y"]), float(attrs["width"]), float(attrs["height"]), ) if all(key in attrs for key in ("left", "top", "right", "bottom")): left = float(attrs["left"]) top = float(attrs["top"]) right = float(attrs["right"]) bottom = float(attrs["bottom"]) return Bounds(left, top, right - left, bottom - top) raw_bounds = attrs.get("bounds") or attrs.get("rect") if isinstance(raw_bounds, dict): return _bounds_from_attrs(raw_bounds) if isinstance(raw_bounds, str): match = _ANDROID_BOUNDS.fullmatch(raw_bounds.strip()) if match: x1 = float(match.group("x1")) y1 = float(match.group("y1")) x2 = float(match.group("x2")) y2 = float(match.group("y2")) return Bounds(x1, y1, x2 - x1, y2 - y1) return None def _text_from_attrs(attrs: dict[str, Any]) -> str | None: for key in ("label", "name", "text", "value", "placeholder"): value = attrs.get(key) if value not in (None, ""): return str(value) return None def _normalize_type(raw_type: str, attrs: dict[str, Any]) -> str: candidate = ( attrs.get("role") or attrs.get("type") or attrs.get("class") or attrs.get("className") or raw_type ) normalized = str(candidate).split(".")[-1].replace("XCUIElementType", "").lower() if "button" in normalized: return "button" if "textfield" in normalized or "textarea" in normalized or "input" in normalized: return "input" if "statictext" in normalized or normalized in {"text", "label"}: return "text" if "image" in normalized or "icon" in normalized: return "image" if "cell" in normalized or "row" in normalized: return "cell" if "window" in normalized: return "window" return normalized or "unknown" def _local_name(tag: str) -> str: return tag.rsplit("}", 1)[-1]