152 lines
4.8 KiB
Python
152 lines
4.8 KiB
Python
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<x1>-?\d+),(?P<y1>-?\d+)\]\[(?P<x2>-?\d+),(?P<y2>-?\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]
|
|
|