cmautobuy/client/src/pdd_device_service.py

"""uiautomator2 设备连接边界。

每次自动化任务通过 ``connect`` 获取一个会话。会话只能在创建它的线程中
使用，退出 ``with`` 后立即释放，避免多个任务同时控制同一台手机。
"""

from __future__ import annotations

import threading
from contextlib import AbstractContextManager
from typing import Any, Callable, Optional

from .performance_timing import current_performance_trace


PDD_PACKAGE_NAME = "com.xunmeng.pinduoduo"


class PddDeviceError(RuntimeError):
    """设备连接失败，并携带供任务状态使用的稳定错误码。"""

    def __init__(self, code: str, message: str) -> None:
        super().__init__(message)
        self.code = code
        self.message = message


class PddDeviceSession(AbstractContextManager[Any]):
    """一个只能由创建线程使用的 uiautomator2 Device 会话。"""

    def __init__(
        self,
        device: Any,
        serial: str,
        initial_app_state: dict[str, Any],
        release: Callable[[], None],
    ) -> None:
        self._device = device
        self.serial = serial
        self.initial_app_state = dict(initial_app_state)
        self._release = release
        self._owner_thread_id = threading.get_ident()
        self._closed = False

    @property
    def device(self) -> Any:
        """返回 Device；跨线程或会话关闭后访问会明确失败。"""

        if self._closed:
            raise PddDeviceError("DEVICE_SESSION_CLOSED", "设备会话已经关闭")
        if threading.get_ident() != self._owner_thread_id:
            raise PddDeviceError(
                "DEVICE_THREAD_VIOLATION",
                "uiautomator2 Device 只能在创建会话的工作线程中使用",
            )
        return self._device

    def __enter__(self) -> Any:
        return self.device

    def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> None:
        if not self._closed:
            self._closed = True
            self._release()
        return None


class PddDeviceService:
    """连接并校验一台已保存的 Android 设备。"""

    def __init__(
        self,
        connector: Optional[Callable[[str], Any]] = None,
    ) -> None:
        self._connector = connector or self._default_connector
        self._state_lock = threading.Lock()
        self._active_serial: Optional[str] = None

    @staticmethod
    def _default_connector(serial: str) -> Any:
        try:
            import uiautomator2 as u2
        except ImportError as exc:
            raise PddDeviceError(
                "DEVICE_U2_MISSING",
                "未安装 uiautomator2，请先安装 client/requirements.txt",
            ) from exc
        return u2.connect(serial)

    @staticmethod
    def _validate_serial(serial: str) -> str:
        value = str(serial or "").strip()
        if not value or any(character.isspace() for character in value):
            raise PddDeviceError("DEVICE_ADDRESS_INVALID", "Android 设备号无效")
        return value

    def connect(self, serial: str) -> PddDeviceSession:
        """连接设备并返回单线程独占会话。

        调用方必须使用 ``with service.connect(serial) as device``。同一个服务
        已有活动会话时会拒绝第二次连接。
        """

        checked_serial = self._validate_serial(serial)
        with self._state_lock:
            if self._active_serial is not None:
                raise PddDeviceError(
                    "DEVICE_IN_USE",
                    f"设备 {self._active_serial} 正在执行其他自动化任务",
                )
            self._active_serial = checked_serial

        try:
            trace = current_performance_trace()
            if trace is None:
                device = self._connector(checked_serial)
            else:
                with trace.stage("uiautomator2_connect"):
                    device = self._connector(checked_serial)
            if trace is None:
                current = device.app_current()
            else:
                with trace.stage("first_app_current"):
                    current = device.app_current()
            if not isinstance(current, dict):
                raise RuntimeError("uiautomator2 未返回有效的设备状态")
        except PddDeviceError:
            self._release(checked_serial)
            raise
        except Exception as exc:
            self._release(checked_serial)
            details = str(exc).lower()
            code = (
                "DEVICE_OFFLINE"
                if any(word in details for word in ("offline", "not found", "disconnected"))
                else "DEVICE_CONNECT_FAILED"
            )
            raise PddDeviceError(
                code,
                f"无法连接 Android 设备 {checked_serial}：{exc}",
            ) from exc

        return PddDeviceSession(
            device,
            checked_serial,
            current,
            lambda: self._release(checked_serial),
        )

    def _release(self, serial: str) -> None:
        with self._state_lock:
            if self._active_serial == serial:
                self._active_serial = None