433 lines
13 KiB
Python
433 lines
13 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""UIAutomator2 client for Android device screen data retrieval.
|
|
|
|
Provides an alternative to the Maestro-based screen API with direct device
|
|
access.
|
|
Handles Maestro blocker detection and removal before connecting.
|
|
"""
|
|
|
|
import base64
|
|
from io import BytesIO
|
|
import re as _re
|
|
import subprocess
|
|
import time
|
|
from typing import TYPE_CHECKING
|
|
import xml.etree.ElementTree as ET
|
|
|
|
from PIL import Image
|
|
from pydantic import BaseModel
|
|
import uiautomator2 as u2
|
|
|
|
from artemis.runtime.awake_service import ensure_device_awake
|
|
from artemis.utils.logger import get_logger
|
|
|
|
if TYPE_CHECKING:
|
|
from uiautomator2 import Device
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
MAESTRO_PACKAGE = "dev.mobile.maestro"
|
|
|
|
|
|
class _CyFunctionDetectorMeta(type):
|
|
def __instancecheck__(self, instance):
|
|
name = type(instance).__name__
|
|
return (
|
|
name
|
|
in (
|
|
"cyfunction",
|
|
"cython_function_or_method",
|
|
"builtin_function_or_method",
|
|
)
|
|
or "cyfunction" in name.lower()
|
|
)
|
|
|
|
|
|
class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta):
|
|
pass
|
|
|
|
|
|
class UIAutomatorScreenData(BaseModel):
|
|
"""Screen data response from UIAutomator2."""
|
|
|
|
model_config = {"ignored_types": (CyFunctionDetector,)}
|
|
|
|
base64: str
|
|
hierarchy_xml: str
|
|
elements: list[dict]
|
|
width: int
|
|
height: int
|
|
|
|
|
|
def _parse_hierarchy_xml_to_elements(hierarchy_xml: str) -> list[dict]:
|
|
"""Parse uiautomator2 XML hierarchy into a flat list of element dictionaries.
|
|
|
|
The output format matches the existing screen API format with attributes
|
|
like:
|
|
- resource-id
|
|
- text
|
|
- content-desc
|
|
- bounds (as string "[x1,y1][x2,y2]")
|
|
- class
|
|
- package
|
|
- checkable, checked, clickable, enabled, focusable, focused
|
|
- scrollable, long-clickable, password, selected
|
|
|
|
Args:
|
|
hierarchy_xml: XML string from uiautomator2.dump_hierarchy()
|
|
|
|
Returns:
|
|
Flat list of element dictionaries
|
|
"""
|
|
elements: list[dict] = []
|
|
|
|
try:
|
|
root = ET.fromstring(hierarchy_xml)
|
|
except ET.ParseError as e:
|
|
logger.error(f"Failed to parse hierarchy XML: {e}")
|
|
return elements
|
|
|
|
def _extract_element(node: ET.Element) -> None:
|
|
"""Recursively extract elements from XML nodes."""
|
|
element: dict = {}
|
|
|
|
for attr_name, attr_value in node.attrib.items():
|
|
if attr_name == "resource-id":
|
|
element["resource-id"] = attr_value
|
|
elif attr_name == "text":
|
|
element["text"] = attr_value
|
|
elif attr_name == "content-desc":
|
|
element["content-desc"] = attr_value
|
|
element["accessibilityText"] = attr_value
|
|
elif attr_name == "bounds":
|
|
element["bounds"] = attr_value
|
|
|
|
match = _re.match(r"\[(\d+),(\d+)\]\[(\d+),(\d+)\]", attr_value)
|
|
if match:
|
|
x1, y1, x2, y2 = map(int, match.groups())
|
|
element["parsed_bounds"] = {
|
|
"left": x1,
|
|
"top": y1,
|
|
"right": x2,
|
|
"bottom": y2,
|
|
}
|
|
elif attr_name == "class":
|
|
element["class"] = attr_value
|
|
elif attr_name == "package":
|
|
element["package"] = attr_value
|
|
elif attr_name in (
|
|
"checkable",
|
|
"checked",
|
|
"clickable",
|
|
"enabled",
|
|
"focusable",
|
|
"focused",
|
|
"scrollable",
|
|
"long-clickable",
|
|
"password",
|
|
"selected",
|
|
):
|
|
element[attr_name] = attr_value
|
|
else:
|
|
element[attr_name] = attr_value
|
|
|
|
if element:
|
|
elements.append(element)
|
|
|
|
for child in node:
|
|
_extract_element(child)
|
|
|
|
_extract_element(root)
|
|
|
|
return elements
|
|
|
|
|
|
def _pil_to_base64(img: Image.Image, format: str = "JPEG", quality: int = 80) -> str:
|
|
"""Convert PIL Image to base64 string with high-quality compression."""
|
|
buffer = BytesIO()
|
|
if format.upper() in ("JPEG", "JPG"):
|
|
format = "JPEG"
|
|
if img.mode in ("RGBA", "P"):
|
|
img = img.convert("RGB")
|
|
img.save(buffer, format=format, quality=quality, optimize=True)
|
|
else:
|
|
img.save(buffer, format=format)
|
|
return base64.b64encode(buffer.getvalue()).decode("utf-8")
|
|
|
|
|
|
def _is_package_installed(device_id: str, pkg: str) -> bool:
|
|
"""Check if a package is installed on the device."""
|
|
try:
|
|
result = subprocess.run(
|
|
["adb", "-s", device_id, "shell", "pm", "list", "packages"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=10,
|
|
)
|
|
if result.returncode != 0:
|
|
logger.warning(f"Failed to list packages: {result.stderr}")
|
|
return False
|
|
|
|
lines = result.stdout.splitlines()
|
|
target = f"package:{pkg}"
|
|
return target in lines
|
|
except subprocess.TimeoutExpired:
|
|
logger.warning("Timeout checking installed packages")
|
|
return False
|
|
except Exception as e:
|
|
logger.warning(f"Error checking installed packages: {e}")
|
|
return False
|
|
|
|
|
|
def _uninstall_package(device_id: str, pkg: str) -> bool:
|
|
"""Uninstall a package from the device for the current user."""
|
|
try:
|
|
result = subprocess.run(
|
|
[
|
|
"adb",
|
|
"-s",
|
|
device_id,
|
|
"shell",
|
|
"pm",
|
|
"uninstall",
|
|
"--user",
|
|
"0",
|
|
pkg,
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=30,
|
|
)
|
|
if result.returncode == 0:
|
|
logger.info(f"Successfully uninstalled {pkg}")
|
|
return True
|
|
else:
|
|
logger.warning(f"Failed to uninstall {pkg}: {result.stderr}")
|
|
return False
|
|
except subprocess.TimeoutExpired:
|
|
logger.warning(f"Timeout uninstalling {pkg}")
|
|
return False
|
|
except Exception as e:
|
|
logger.warning(f"Error uninstalling {pkg}: {e}")
|
|
return False
|
|
|
|
|
|
def _ensure_maestro_not_installed(device_id: str) -> None:
|
|
"""Check if Maestro is installed and uninstall it.
|
|
|
|
Maestro conflicts with uiautomator2 - if Maestro's package is installed,
|
|
u2.connect() will fail. This function ensures Maestro is removed before
|
|
attempting to connect.
|
|
"""
|
|
if _is_package_installed(device_id, MAESTRO_PACKAGE):
|
|
logger.warning(
|
|
f"Maestro ({MAESTRO_PACKAGE}) detected - uninstalling to enable UIAutomator2..."
|
|
)
|
|
_uninstall_package(device_id, MAESTRO_PACKAGE)
|
|
|
|
|
|
class UIAutomatorClient:
|
|
"""UIAutomator2 client for Android screen data retrieval.
|
|
|
|
This client uses uiautomator2 library for direct device communication,
|
|
providing faster hierarchy and screenshot retrieval compared to Maestro.
|
|
|
|
Important: Maestro must not be installed on the device as it conflicts
|
|
with uiautomator2. This client automatically handles Maestro removal.
|
|
"""
|
|
|
|
def __init__(self, device_id: str):
|
|
"""Initialize the UIAutomator client.
|
|
|
|
Args:
|
|
device_id: The Android device serial number (e.g., "emulator-5554")
|
|
"""
|
|
self._device_id = device_id
|
|
self._device: Device | None = None
|
|
self._awake_strategy: str | None = None
|
|
|
|
def _ensure_connected(self) -> "Device":
|
|
"""Ensure connection to the device, handling Maestro blocker.
|
|
|
|
Returns:
|
|
Connected uiautomator2 Device instance
|
|
"""
|
|
if self._device is not None:
|
|
try:
|
|
# Quick check if connection is still alive
|
|
self._device.info
|
|
return self._device
|
|
except Exception:
|
|
logger.warning("UIAutomator2 connection lost, reconnecting...")
|
|
self._device = None
|
|
|
|
# Ensure Maestro is not blocking us
|
|
_ensure_maestro_not_installed(self._device_id)
|
|
|
|
# Enroll the device in the containing Artemis service lifetime before
|
|
# UIAutomator2 reads the screen. Client disconnects do not release it.
|
|
if self._awake_strategy is None:
|
|
self._awake_strategy = ensure_device_awake(self._device_id)
|
|
|
|
# Connect to device
|
|
logger.info(f"Connecting UIAutomator2 to device: {self._device_id}")
|
|
try:
|
|
self._device = u2.connect(self._device_id)
|
|
except Exception:
|
|
self._awake_strategy = None
|
|
raise
|
|
logger.info("UIAutomator2 connected successfully")
|
|
return self._device
|
|
|
|
def connect(self) -> None:
|
|
"""Connect and activate this session's awake strategy without an action."""
|
|
self._ensure_connected()
|
|
|
|
def press_key(self, key: str):
|
|
"""Press a key on the device.
|
|
|
|
Args:
|
|
key: Key to press (e.g., "home", "back", "enter"...)
|
|
"""
|
|
device = self._ensure_connected()
|
|
return device.press(key=key)
|
|
|
|
def send_text(self, text: str) -> None:
|
|
"""Send text input to the device using FastInputIME.
|
|
|
|
This method supports special characters (e.g., 'ö') that ADB shell
|
|
input text cannot handle. The original IME is restored after input.
|
|
|
|
Args:
|
|
text: The text to input
|
|
"""
|
|
|
|
device = self._ensure_connected()
|
|
device.set_fastinput_ime(True)
|
|
# Allow IME to initialize
|
|
time.sleep(0.3)
|
|
try:
|
|
device.send_keys(text)
|
|
# Give FastInputIME time to process the broadcast and commit text
|
|
# before switching it off and killing it.
|
|
time.sleep(0.5)
|
|
finally:
|
|
device.set_fastinput_ime(False)
|
|
|
|
def clear_text(self) -> bool:
|
|
"""Clear text of the currently focused element using UIAutomator2.
|
|
|
|
Returns True on success, False on failure.
|
|
"""
|
|
try:
|
|
device = self._ensure_connected()
|
|
device(focused=True).clear_text()
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"UIAutomator2 clear_text failed: {e}")
|
|
return False
|
|
|
|
def get_hierarchy(self) -> str:
|
|
"""Get the UI hierarchy XML from the device.
|
|
|
|
Returns:
|
|
Compressed UI hierarchy XML string
|
|
"""
|
|
device = self._ensure_connected()
|
|
return device.dump_hierarchy(compressed=True)
|
|
|
|
def get_screenshot(self) -> Image.Image | None:
|
|
"""Capture a screenshot from the device.
|
|
|
|
Uses adb screencap for more reliable capture of all layers (including
|
|
popups), falling back to UIAutomator2 if it fails.
|
|
|
|
Returns:
|
|
PIL Image or None if capture failed
|
|
"""
|
|
try:
|
|
result = subprocess.run(
|
|
["adb", "-s", self._device_id, "exec-out", "screencap", "-p"],
|
|
capture_output=True,
|
|
timeout=10,
|
|
check=True,
|
|
)
|
|
return Image.open(BytesIO(result.stdout))
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Failed to capture screenshot via adb: {e}, falling back to uiautomator2"
|
|
)
|
|
device = self._ensure_connected()
|
|
return device.screenshot()
|
|
|
|
def get_screenshot_base64(self) -> str | None:
|
|
"""Capture a screenshot and return as base64 string.
|
|
|
|
Returns:
|
|
Base64 encoded PNG screenshot or None if capture failed
|
|
"""
|
|
screenshot = self.get_screenshot()
|
|
if screenshot is None:
|
|
return None
|
|
return _pil_to_base64(screenshot, format="JPEG")
|
|
|
|
def get_screen_data(self) -> UIAutomatorScreenData:
|
|
"""Get complete screen data including screenshot and hierarchy.
|
|
|
|
Returns:
|
|
UIAutomatorScreenData with screenshot, hierarchy, elements, and
|
|
dimensions
|
|
"""
|
|
device = self._ensure_connected()
|
|
|
|
# Get screenshot
|
|
screenshot = self.get_screenshot()
|
|
if screenshot is None:
|
|
raise RuntimeError("Failed to capture screenshot")
|
|
|
|
# Get hierarchy
|
|
hierarchy_xml = device.dump_hierarchy(compressed=True)
|
|
|
|
# Parse XML to flat elements list
|
|
elements = _parse_hierarchy_xml_to_elements(hierarchy_xml)
|
|
|
|
return UIAutomatorScreenData(
|
|
base64=_pil_to_base64(screenshot, format="JPEG"),
|
|
hierarchy_xml=hierarchy_xml,
|
|
elements=elements,
|
|
width=screenshot.width,
|
|
height=screenshot.height,
|
|
)
|
|
|
|
def disconnect(self) -> None:
|
|
"""Disconnect this client without ending the host's awake lifetime."""
|
|
self._awake_strategy = None
|
|
self._device = None
|
|
logger.info("UIAutomator2 client disconnected")
|
|
|
|
|
|
def get_client(device_id: str) -> UIAutomatorClient:
|
|
"""Factory function to create a UIAutomatorClient.
|
|
|
|
Args:
|
|
device_id: The Android device serial number
|
|
|
|
Returns:
|
|
UIAutomatorClient instance
|
|
"""
|
|
return UIAutomatorClient(device_id=device_id)
|