fix takeScreenshot for display_id, change double flash launch to one …

…flash

uiautomator2/__init__.py

@@ -6,6 +6,7 @@
 import base64
 import contextlib
 import dataclasses
+import io
 import logging
 import os
 import re
@@ -18,6 +19,7 @@
 import adbutils
 from lxml import etree
 from retry import retry
+from PIL import Image
 
 from uiautomator2.core import BasicUiautomatorServer
 
@@ -28,7 +30,7 @@
 from uiautomator2.exceptions import AdbShellError, BaseException, ConnectError, DeviceError, HierarchyEmptyError, SessionBrokenError
 from uiautomator2.settings import Settings
 from uiautomator2.swipe import SwipeExt
-from uiautomator2.utils import list2cmdline, deprecated
+from uiautomator2.utils import image_convert, list2cmdline, deprecated
 from uiautomator2.watcher import WatchContext, Watcher
 from uiautomator2.abstract import AbstractShell, AbstractUiautomatorServer, ShellResponse
 
@@ -253,19 +255,17 @@ def screenshot(self, filename: Optional[str] = None, format="pillow", display_id
             screenshot().save("saved.png")
             cv2.imwrite('saved.jpg', screenshot(format='opencv'))
         """
-        pil_img = self._dev.screenshot(display_id=display_id)
+        if display_id is None:
+            base64_data = self.jsonrpc.takeScreenshot(1, 80)
+            jpg_raw = base64.b64decode(base64_data)
+            pil_img = Image.open(io.BytesIO(jpg_raw))
+        else:
+            pil_img = self._dev.screenshot(display_id=display_id)
+
         if filename:
             pil_img.save(filename)
             return
-        if format == 'pillow':
-            return pil_img
-        elif format == 'opencv':
-            pil_img = pil_img.convert("RGB")
-            import cv2
-            import numpy as np
-            return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
-        elif format == 'raw':
-            return pil_img.tobytes()
+        return image_convert(pil_img, format)
 
     def dump_hierarchy(self, compressed=False, pretty=False, max_depth: int = None) -> str:
         """

uiautomator2/core.py

@@ -138,6 +138,8 @@ def _jsonrpc_call(dev: adbutils.AdbDevice, method: str, params: Any, timeout: fl
         if "android.os.DeadObjectException" in message:
             # https://developer.android.com/reference/android/os/DeadObjectException
             raise UiAutomationNotConnectedError("android.os.DeadObjectException")
+        if "android.os.DeadSystemRuntimeException":
+            raise UiAutomationNotConnectedError("android.os.DeadSystemRuntimeException")
         if "uiautomator.UiObjectNotFoundException" in message:
             raise UiObjectNotFoundError(code, message, params)
         if "java.lang.StackOverflowError" in message:
@@ -158,7 +160,7 @@ def __init__(self, dev: adbutils.AdbDevice) -> None:
         self._process = None
         self._lock = threading.Lock()
         self._debug = False
-        self.start_uiautomator()
+        self.start_uiautomator(_silent=True)
         atexit.register(self.stop_uiautomator, wait=False)
 
     @property
@@ -169,16 +171,16 @@ def debug(self) -> bool:
     def debug(self, value: bool):
         self._debug = bool(value)
 
-    def start_uiautomator(self):
+    def start_uiautomator(self, _silent: bool = False):
         try:
-            self._do_start_uiautomator()
+            self._do_start_uiautomator(silent=_silent)
         except APKSignatureError as e:
             logger.debug("APkSignatureError: %s", e)
             self._dev.uninstall("com.github.uiautomator")
             self._dev.uninstall("com.github.uiautomator.test")
-            self._do_start_uiautomator()
+            self._do_start_uiautomator(silent=_silent)
 
-    def _do_start_uiautomator(self):
+    def _do_start_uiautomator(self, silent: bool):
         """
         Start uiautomator2 server
 
@@ -192,7 +194,7 @@ def _do_start_uiautomator(self):
                     self._process = None
             if not self._check_alive():
                 self._process = launch_uiautomator(self._dev)
-                self._wait_ready()
+                self._wait_ready(show_float_window=not silent)
 
     def _setup_apks(self):
         assets_dir = Path(__file__).parent / "assets"
@@ -230,6 +232,7 @@ def _install_apk(self, apk_path: Path):
         self._dev.shell(['pm', 'install', '-r', '-t', '-d', target_path])
 
     def _wait_instrument_ready(self, timeout: float):
+        """wait until "INSTRUMENTATION_STATUS_CODE: 1" show up"""
         deadline = time.time() + timeout
         while time.time() < deadline:
             output = self._process.output.decode("utf-8", errors="ignore")
@@ -268,14 +271,15 @@ def _check_alive(self) -> bool:
         except HTTPError:
             return False
 
-    def _wait_ready(self, launch_timeout=30, service_timeout=30):
+    def _wait_ready(self, launch_timeout=30, service_timeout=30, show_float_window: bool = True):
         """Wait until uiautomator2 server is ready"""
         # wait am instrument start
         self._wait_instrument_ready(launch_timeout)
-        # launch a toast window to make sure uiautomator is alive
-        logger.debug("show float window")
         self._dev.shell("am startservice -a com.github.uiautomator.ACTION_START")
-        self._dev.shell("am start -n com.github.uiautomator/.ToastActivity -e showFloatWindow true")
+        if show_float_window:
+            # launch a toast window to make sure uiautomator is alive
+            logger.debug("show float window")
+            self._dev.shell("am start -n com.github.uiautomator/.ToastActivity -e showFloatWindow true")
         self._wait_stub_ready(service_timeout)
         time.sleep(1) # wait ATX goto background
 

uiautomator2/exceptions.py

@@ -4,10 +4,10 @@
 class BaseException(Exception):
     """ base error for uiautomator2 """
 
-## DeviceError
-class DeviceError(BaseException):
-    pass
+class MissingLibError(BaseException): ...
 
+## DeviceError
+class DeviceError(BaseException): ...
 class AdbShellError(DeviceError):...
 class ConnectError(DeviceError):...
 class HTTPError(DeviceError):...

uiautomator2/utils.py

@@ -8,9 +8,10 @@
 import typing
 from typing import Union
 import warnings
+from PIL import Image
 
 from uiautomator2._proto import Direction
-from uiautomator2.exceptions import SessionBrokenError, UiObjectNotFoundError
+from uiautomator2.exceptions import MissingLibError, SessionBrokenError, UiObjectNotFoundError
 
 
 def check_alive(fn):
@@ -246,6 +247,22 @@ def wrapper(*args, **kwargs):
     return decorator
 
 
+def image_convert(im: Image.Image, format: str):
+    if format == "pillow":
+        return im
+    if format == "opencv":
+        try:
+            import cv2
+            import numpy as np
+            im = im.convert("RGB")
+            return cv2.cvtColor(np.array(im), cv2.COLOR_RGB2BGR)
+        except ImportError:
+            raise MissingLibError("missing lib: cv2 or numpy")
+    if format == "raw":
+        return im.tobytes()
+    raise ValueError("Unsupported format:", format)
+
+
 if __name__ == "__main__":
     for n in (1, 10000, 10000000, 10000000000):
         print(n, natualsize(n))