Skip to content

Commit

Permalink
fix android dumphierarchy and ios screenshot
Browse files Browse the repository at this point in the history
  • Loading branch information
codeskyblue committed May 25, 2024
1 parent 01c5c16 commit 32a6d25
Show file tree
Hide file tree
Showing 4 changed files with 41 additions and 39 deletions.
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ httpx = "*"
fastapi = "^0.111.0"
uvicorn = {version = "*", extras = ["standard"]}
poetry = "^1.8.2"
pydantic = "^2.6"

[tool.poetry.extras]
appium = ["appium-python-client", "httppretty"]
Expand Down
59 changes: 24 additions & 35 deletions uiautodev/driver/android.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
import re
import time
from functools import cached_property, partial
from typing import List, Tuple
from typing import List, Optional, Tuple
from xml.etree import ElementTree

import adbutils
Expand All @@ -20,7 +20,7 @@
from uiautodev.driver.base_driver import BaseDriver
from uiautodev.driver.udt.udt import UDT, UDTError
from uiautodev.exceptions import AndroidDriverException, RequestError
from uiautodev.model import Node, ShellResponse, WindowSize
from uiautodev.model import Node, Rect, ShellResponse, WindowSize
from uiautodev.utils.common import fetch_through_socket

logger = logging.getLogger(__name__)
Expand All @@ -31,8 +31,8 @@ def __init__(self, serial: str):
self.adb_device = adbutils.device(serial)
self._try_dump_list = [
self._get_u2_hierarchy,
self._get_appium_hierarchy,
self._get_udt_dump_hierarchy,
# self._get_appium_hierarchy,
]

@cached_property
Expand Down Expand Up @@ -64,16 +64,16 @@ def shell(self, command: str) -> ShellResponse:
except Exception as e:
return ShellResponse(output="", error=f"adb error: {str(e)}")

def dump_hierarchy(self) -> Tuple[str, Node]:
def dump_hierarchy(self, display_id: Optional[int] = 0) -> Tuple[str, Node]:
"""returns xml string and hierarchy object"""
wsize = self.adb_device.window_size()
logger.debug("window size: %s", wsize)
start = time.time()
xml_data = self._dump_hierarchy_raw()
logger.debug("dump_hierarchy cost: %s", time.time() - start)

wsize = self.adb_device.window_size()
logger.debug("window size: %s", wsize)
return xml_data, parse_xml(
xml_data, WindowSize(width=wsize[0], height=wsize[1])
xml_data, WindowSize(width=wsize[0], height=wsize[1]), display_id
)

def _dump_hierarchy_raw(self) -> str:
Expand All @@ -99,28 +99,6 @@ def _dump_hierarchy_raw(self) -> str:
def _get_u2_hierarchy(self) -> str:
d = u2.connect_usb(self.serial)
return d.dump_hierarchy()
# c = self.device.create_connection(adbutils.Network.TCP, 9008)
# try:
# compressed = False
# payload = {
# "jsonrpc": "2.0",
# "method": "dumpWindowHierarchy",
# "params": [compressed],
# "id": 1,
# }
# content = fetch_through_socket(
# c, "/jsonrpc/0", method="POST", json=payload, timeout=5
# )
# json_resp = json.loads(content)
# if "error" in json_resp:
# raise AndroidDriverException(json_resp["error"])
# return json_resp["result"]
# except adbutils.AdbError as e:
# raise AndroidDriverException(
# f"Failed to get hierarchy from u2 server: {str(e)}"
# )
# finally:
# c.close()

def _get_appium_hierarchy(self) -> str:
c = self.adb_device.create_connection(adbutils.Network.TCP, 6790)
Expand Down Expand Up @@ -168,43 +146,54 @@ def wake_up(self):
self.adb_device.keyevent("WAKEUP")


def parse_xml(xml_data: str, wsize: WindowSize) -> Node:
def parse_xml(xml_data: str, wsize: WindowSize, display_id: int) -> Node:
root = ElementTree.fromstring(xml_data)
return parse_xml_element(root, wsize)
node = parse_xml_element(root, wsize, display_id)
if node is None:
raise AndroidDriverException("Failed to parse xml")
return node


def parse_xml_element(
element, wsize: WindowSize, indexes: List[int] = [0]
) -> Node:
def parse_xml_element(element, wsize: WindowSize, display_id: int, indexes: List[int] = [0]) -> Optional[Node]:
"""
Recursively parse an XML element into a dictionary format.
"""
name = element.tag
if name == "node":
name = element.attrib.get("class", "node")
curr_display_id = int(element.attrib.get("display-id", display_id))
if curr_display_id != display_id:
return

bounds = None
rect = None
# eg: bounds="[883,2222][1008,2265]"
if "bounds" in element.attrib:
bounds = element.attrib["bounds"]
bounds = list(map(int, re.findall(r"\d+", bounds)))
assert len(bounds) == 4
rect = Rect(x=bounds[0], y=bounds[1], width=bounds[2] - bounds[0], height=bounds[3] - bounds[1])
bounds = (
bounds[0] / wsize.width,
bounds[1] / wsize.height,
bounds[2] / wsize.width,
bounds[3] / wsize.height,
)
bounds = map(partial(round, ndigits=4), bounds)

elem = Node(
key="-".join(map(str, indexes)),
name=name,
bounds=bounds,
rect=rect,
properties={key: element.attrib[key] for key in element.attrib},
children=[],
)

# Construct xpath for children
for index, child in enumerate(element):
elem.children.append(parse_xml_element(child, wsize, indexes + [index]))
child_node = parse_xml_element(child, wsize, display_id, indexes + [index])
if child_node:
elem.children.append(child_node)

return elem
8 changes: 8 additions & 0 deletions uiautodev/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,18 @@ class ShellResponse(BaseModel):
error: Optional[str] = ""


class Rect(BaseModel):
x: int
y: int
width: int
height: int


class Node(BaseModel):
key: str
name: str
bounds: Optional[Tuple[float, float, float, float]] = None
rect: Optional[Rect] = None
properties: Dict[str, Union[str, bool]] = []
children: List[Node] = []

Expand Down
12 changes: 8 additions & 4 deletions uiautodev/router/device.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
from uiautodev.model import DeviceInfo, Node, ShellResponse
from uiautodev.provider import BaseProvider


logger = logging.getLogger(__name__)

class AndroidShellPayload(BaseModel):
Expand Down Expand Up @@ -58,7 +57,7 @@ def _screenshot(serial: str, id: int) -> Response:
"""Take a screenshot of device"""
try:
driver = provider.get_device_driver(serial)
pil_img = driver.screenshot(id)
pil_img = driver.screenshot(id).convert("RGB")
buf = io.BytesIO()
pil_img.save(buf, format="JPEG")
image_bytes = buf.getvalue()
Expand All @@ -68,12 +67,17 @@ def _screenshot(serial: str, id: int) -> Response:
return Response(content=str(e), media_type="text/plain", status_code=500)

@router.get("/{serial}/hierarchy")
def dump_hierarchy(serial: str) -> Node:
def dump_hierarchy(serial: str, format: str = "json") -> Node:
"""Dump the view hierarchy of an Android device"""
try:
driver = provider.get_device_driver(serial)
xml_data, hierarchy = driver.dump_hierarchy()
return hierarchy
if format == "xml":
return Response(content=xml_data, media_type="text/xml")
elif format == "json":
return hierarchy
else:
return Response(content=f"Invalid format: {format}", media_type="text/plain", status_code=400)
except Exception as e:
logger.exception("dump_hierarchy failed")
return Response(content=str(e), media_type="text/plain", status_code=500)
Expand Down

0 comments on commit 32a6d25

Please sign in to comment.