android-env 1.2.2__py3-none-any.whl → 1.2.3__py3-none-any.whl

android_env/components/adb_call_parser.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -20,7 +20,6 @@ import re
 import subprocess
 import sys
 import tempfile
-from typing import Optional
 
 from absl import logging
 from android_env.components import adb_controller as adb_control
@@ -41,9 +40,8 @@ _BUTTON_TO_KEYCODE = {
 class AdbCallParser:
   """Parses AdbRequest messages and executes corresponding adb commands."""
 
-  def __init__(self, adb_controller: adb_control.AdbController, tmp_dir: str):
+  def __init__(self, adb_controller: adb_control.AdbController):
     self._adb_controller = adb_controller
-    self._tmp_dir = tmp_dir
     self._handlers = {
         'install_apk': self._install_apk,
         'start_activity': self._start_activity,
@@ -65,7 +63,7 @@ class AdbCallParser:
     }
 
   def _execute_command(
-      self, command_args: list[str], timeout: Optional[float]
+      self, command_args: list[str], timeout: float | None
   ) -> tuple[adb_pb2.AdbResponse, bytes]:
     """Executes the command, catches errors and populates the response status.
 
@@ -109,12 +107,12 @@ class AdbCallParser:
                                 f'Got: {request.timeout_sec}')
       return response
 
-    timeout: Optional[float] = request.timeout_sec or None
+    timeout: float | None = request.timeout_sec or None
     return self._handlers[command_type](request, timeout)
 
-  def _force_stop(self,
-                  request: adb_pb2.AdbRequest,
-                  timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+  def _force_stop(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Stops an application.
 
     Args:
@@ -137,9 +135,9 @@ class AdbCallParser:
 
     return response
 
-  def _fetch_current_task_id(self,
-                             full_activity_name: str,
-                             timeout: Optional[float] = None) -> int:
+  def _fetch_current_task_id(
+      self, full_activity_name: str, timeout: float | None = None
+  ) -> int:
     """Returns the task ID of the given `full_activity_name`.
 
     Args:
@@ -185,9 +183,8 @@ class AdbCallParser:
     return -1
 
   def _start_screen_pinning(
-      self,
-      request: adb_pb2.AdbRequest,
-      timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Pins an application.
 
     Args:
@@ -219,9 +216,9 @@ class AdbCallParser:
 
     return response
 
-  def _send_broadcast(self,
-                      request: adb_pb2.AdbRequest,
-                      timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+  def _send_broadcast(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Sends a broadcast.
 
     Args:
@@ -232,21 +229,29 @@ class AdbCallParser:
       An AdbResponse.
     """
 
-    send_brodcast = request.send_broadcast
+    send_broadcast = request.send_broadcast
     response = adb_pb2.AdbResponse(status=adb_pb2.AdbResponse.Status.OK)
-    if not send_brodcast.action:
+    if not send_broadcast.action:
       response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
       response.error_message = ('`send_broadcast.{action}` cannot be empty.')
       return response
 
+    if send_broadcast.component:
+      component_args = ['-n', send_broadcast.component]
+    else:
+      component_args = []
+
     response, _ = self._execute_command(
-        ['shell', 'am', 'broadcast', '-a', send_brodcast.action],
-        timeout=timeout)
+        ['shell', 'am', 'broadcast', '-a', send_broadcast.action]
+        + component_args,
+        timeout=timeout,
+    )
+
     return response
 
-  def _install_apk(self,
-                   request: adb_pb2.AdbRequest,
-                   timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+  def _install_apk(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Installs an app given its local path in the filesystem.
 
     Args:
@@ -262,25 +267,46 @@ class AdbCallParser:
     response = adb_pb2.AdbResponse()
     location_type = install_apk.WhichOneof('location')
     logging.info('location_type: %s', location_type)
-    if location_type != 'filesystem':
-      response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
-      response.error_message = (
-          f'Unsupported `install_apk.location` type: {location_type}')
-      return response
 
-    fpath = install_apk.filesystem.path
-    if not os.path.exists(fpath):
-      response.status = adb_pb2.AdbResponse.Status.INTERNAL_ERROR
-      response.error_message = f'Could not find local_apk_path: {fpath}'
-      return response
+    match location_type:
+      case 'filesystem':
+        fpath = install_apk.filesystem.path
+        if not os.path.exists(fpath):
+          response.status = adb_pb2.AdbResponse.Status.INTERNAL_ERROR
+          response.error_message = f'Could not find local_apk_path: {fpath}'
+          return response
+
+        response, _ = self._execute_command(
+            ['install', '-r', '-t', '-g', fpath], timeout=timeout
+        )
+      case 'blob':
+
+        # `delete_on_close` was only added in Python 3.12 so we add a switch
+        # here to still support previous Python versions.
+        if sys.version_info >= (3, 12):
+          kwargs = {'suffix': '.apk', 'delete_on_close': False}
+        else:
+          kwargs = {'suffix': '.apk'}
+
+        with tempfile.NamedTemporaryFile(**kwargs) as f:
+          fpath = f.name
+          f.write(install_apk.blob.contents)
+
+          response, _ = self._execute_command(
+              ['install', '-r', '-t', '-g', fpath], timeout=timeout
+          )
+      case _:
+        response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
+        response.error_message = (
+            f'Unsupported `install_apk.location` type: {location_type}'
+        )
+        return response
 
-    response, _ = self._execute_command(['install', '-r', '-t', '-g', fpath],
-                                        timeout=timeout)
     return response
 
-  def _start_activity(self,
-                      request: adb_pb2.AdbRequest,
-                      timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+  def _start_activity(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Starts a given activity.
 
     Options for `start_activity`:
@@ -326,9 +352,9 @@ class AdbCallParser:
     response.start_activity.output = command_output
     return response
 
-  def _press_button(self,
-                    request: adb_pb2.AdbRequest,
-                    timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+  def _press_button(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Presses a keyboard key.
 
     Args:
@@ -354,9 +380,8 @@ class AdbCallParser:
     return response
 
   def _handle_uninstall_package(
-      self,
-      request: adb_pb2.AdbRequest,
-      timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Handles UninstallPackage messages.
 
     Args:
@@ -394,9 +419,8 @@ class AdbCallParser:
     return response
 
   def _get_current_activity(
-      self,
-      request: adb_pb2.AdbRequest,
-      timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Fetches current activity.
 
     Args:
@@ -448,9 +472,9 @@ class AdbCallParser:
     response.get_current_activity.full_activity = matches.group(1)
     return response
 
-  def _get_orientation(self,
-                       request: adb_pb2.AdbRequest,
-                       timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+  def _get_orientation(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Fetches current device orientation.
 
     Args:
@@ -486,23 +510,30 @@ class AdbCallParser:
       if physical_width:
         skip_next = int(physical_width.group(1)) < 0
 
-      surface_orientation = re.match(r'\s+SurfaceOrientation:\s+(\d)', line)
+      surface_orientation = re.match(
+          r'\s+(SurfaceOrientation|InputDeviceOrientation):\s+(\d)', line
+      )
+
       if surface_orientation is not None:
         if skip_next:
           continue
-        orientation = surface_orientation.group(1)
+        if surface_orientation.re.groups < 2:
+          continue
+        orientation = surface_orientation.group(2)
         logging.info('Done getting orientation: %r', orientation)
         response.get_orientation.orientation = int(orientation)
         return response
 
     response.status = adb_pb2.AdbResponse.Status.INTERNAL_ERROR
-    response.error_message = ('Could not find SurfaceOrientation in dumpsys '
-                              'output')
+    response.error_message = (
+        'Could not find SurfaceOrientation/InputDeviceOrientation in dumpsys '
+        'output'
+    )
     return response
 
-  def _push(self,
-            request: adb_pb2.AdbRequest,
-            timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+  def _push(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Uploads contents to the device.
 
     Args:
@@ -520,7 +551,7 @@ class AdbCallParser:
           error_message='Push.path is empty.')
 
     # Create temporary file with `push` contents.
-    with tempfile.NamedTemporaryFile(dir=self._tmp_dir, delete=False) as f:
+    with tempfile.NamedTemporaryFile(delete=False) as f:
       fname = f.name
       f.write(request.push.content)
     # Issue `adb push` command to upload file.
@@ -531,9 +562,9 @@ class AdbCallParser:
 
     return response
 
-  def _pull(self,
-            request: adb_pb2.AdbRequest,
-            timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+  def _pull(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Downloads file content from the device.
 
     Args:
@@ -551,7 +582,7 @@ class AdbCallParser:
           error_message='Pull.path is empty.')
 
     # Issue `adb pull` command to copy it to a temporary file.
-    with tempfile.NamedTemporaryFile(dir=self._tmp_dir, delete=False) as f:
+    with tempfile.NamedTemporaryFile(delete=False) as f:
       fname = f.name
       logging.info('Downloading %r to %r.', path, fname)
       response, _ = self._execute_command(['pull', path, fname],
@@ -564,9 +595,9 @@ class AdbCallParser:
 
     return response
 
-  def _input_text(self,
-                  request: adb_pb2.AdbRequest,
-                  timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+  def _input_text(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Inserts text as keyboard events.
 
     Args:
@@ -587,9 +618,9 @@ class AdbCallParser:
                                         timeout=timeout)
     return response
 
-  def _tap(self,
-           request: adb_pb2.AdbRequest,
-           timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+  def _tap(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Taps the device screen.
 
     Args:
@@ -617,9 +648,9 @@ class AdbCallParser:
 
     return response
 
-  def _handle_settings(self,
-                       request: adb_pb2.AdbRequest,
-                       timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+  def _handle_settings(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Handles SettingsRequest messages.
 
     Args:
@@ -642,70 +673,84 @@ class AdbCallParser:
     namespace = adb_pb2.AdbRequest.SettingsRequest.Namespace.Name(
         request.name_space).lower()
 
-    verb = request.WhichOneof('verb')
-    if verb == 'get':
-      get = request.get
-      if not get.key:
-        response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
-        response.error_message = (
-            f'Empty SettingsRequest.get.key. Got: {request}.')
-        return response
-      response, command_output = self._execute_command(
-          ['shell', 'settings', 'get', namespace, get.key], timeout=timeout)
-      response.settings.output = command_output
-    elif verb == 'put':
-      put = request.put
-      if not put.key or not put.value:
-        response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
-        response.error_message = (
-            f'Empty SettingsRequest.put key or value. Got: {request}.')
-        return response
-      response, command_output = self._execute_command(
-          ['shell', 'settings', 'put', namespace, put.key, put.value],
-          timeout=timeout)
-      response.settings.output = command_output
-    elif verb == 'delete_key':
-      delete = request.delete_key
-      if not delete.key:
+    match request.WhichOneof('verb'):
+      case 'get':
+        get = request.get
+        if not get.key:
+          response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
+          response.error_message = (
+              f'Empty SettingsRequest.get.key. Got: {request}.'
+          )
+          return response
+        response, command_output = self._execute_command(
+            ['shell', 'settings', 'get', namespace, get.key], timeout=timeout
+        )
+        response.settings.output = command_output
+      case 'put':
+        put = request.put
+        if not put.key or not put.value:
+          response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
+          response.error_message = (
+              f'Empty SettingsRequest.put key or value. Got: {request}.'
+          )
+          return response
+        response, command_output = self._execute_command(
+            ['shell', 'settings', 'put', namespace, put.key, put.value],
+            timeout=timeout,
+        )
+        response.settings.output = command_output
+      case 'delete_key':
+        delete = request.delete_key
+        if not delete.key:
+          response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
+          response.error_message = (
+              f'Empty SettingsRequest.delete_key.key. Got: {request}.'
+          )
+          return response
+        response, command_output = self._execute_command(
+            ['shell', 'settings', 'delete', namespace, delete.key],
+            timeout=timeout,
+        )
+        response.settings.output = command_output
+      case 'reset':
+        reset = request.reset
+        # At least one of `package_name` or `mode` should be given.
+        if (
+            not reset.package_name
+            and reset.mode
+            == adb_pb2.AdbRequest.SettingsRequest.Reset.Mode.UNKNOWN
+        ):
+          response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
+          response.error_message = (
+              'At least one of SettingsRequest.reset package_name or mode'
+              f' should be given. Got: {request}.'
+          )
+          return response
+
+        mode = adb_pb2.AdbRequest.SettingsRequest.Reset.Mode.Name(
+            reset.mode
+        ).lower()
+        arg = reset.package_name or mode
+        response, command_output = self._execute_command(
+            ['shell', 'settings', 'reset', namespace, arg], timeout=timeout
+        )
+        response.settings.output = command_output
+      case 'list':
+        response, command_output = self._execute_command(
+            ['shell', 'settings', 'list', namespace], timeout=timeout
+        )
+        response.settings.output = command_output
+      case _:
         response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
         response.error_message = (
-            f'Empty SettingsRequest.delete_key.key. Got: {request}.')
-        return response
-      response, command_output = self._execute_command(
-          ['shell', 'settings', 'delete', namespace, delete.key],
-          timeout=timeout)
-      response.settings.output = command_output
-    elif verb == 'reset':
-      reset = request.reset
-      # At least one of `package_name` or `mode` should be given.
-      if (not reset.package_name and
-          reset.mode == adb_pb2.AdbRequest.SettingsRequest.Reset.Mode.UNKNOWN):
-        response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
-        response.error_message = (
-            'At least one of SettingsRequest.reset package_name or mode should '
-            f'be given. Got: {request}.')
-        return response
-
-      mode = adb_pb2.AdbRequest.SettingsRequest.Reset.Mode.Name(
-          reset.mode).lower()
-      arg = reset.package_name or mode
-      response, command_output = self._execute_command(
-          ['shell', 'settings', 'reset', namespace, arg], timeout=timeout)
-      response.settings.output = command_output
-    elif verb == 'list':
-      response, command_output = self._execute_command(
-          ['shell', 'settings', 'list', namespace], timeout=timeout)
-      response.settings.output = command_output
-    else:
-      response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
-      response.error_message = (
-          f'Unknown SettingsRequest.verb. Got: {request}.')
+            f'Unknown SettingsRequest.verb. Got: {request}.'
+        )
 
     return response
 
-  def _handle_generic(self,
-                      request: adb_pb2.AdbRequest,
-                      timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+  def _handle_generic(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Handles GenericRequest messages.
 
     Args:
@@ -723,9 +768,8 @@ class AdbCallParser:
     return response
 
   def _handle_package_manager(
-      self,
-      request: adb_pb2.AdbRequest,
-      timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Handles PackageManagerRequest messages.
 
     Args:
@@ -740,60 +784,64 @@ class AdbCallParser:
     request = request.package_manager
     response = adb_pb2.AdbResponse()
 
-    verb = request.WhichOneof('verb')
-    if verb == 'list':
-      what = request.list.WhichOneof('what')
-      response, output = self._execute_command(['shell', 'pm', 'list', what],
-                                               timeout=timeout)
-
-      if output:
-        items = output.decode('utf-8').split()
-        # Remove prefix for each item.
-        prefix = {
-            'features': 'feature:',
-            'libraries': 'library:',
-            'packages': 'package:',
-        }[what]
-        items = [x[len(prefix):] for x in items if x.startswith(prefix)]
-        response.package_manager.list.items.extend(items)
-      response.package_manager.output = output
-    elif verb == 'clear':
-      package_name = request.clear.package_name
-      if not package_name:
-        response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
-        response.error_message = (
-            f'Empty PackageManagerRequest.clear.package_name. Got: {request}.')
-        return response
-
-      args = ['shell', 'pm', 'clear', package_name]
-      if request.clear.user_id:
-        args.insert(3, '-f')
-        args.insert(4, request.clear.user_id)
-      response, response.package_manager.output = self._execute_command(
-          args, timeout=timeout)
-    elif verb == 'grant':
-      grant = request.grant
-      if not grant.package_name:
-        response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
-        response.error_message = ('`grant.package_name` cannot be empty.')
-        return response
-
-      if not grant.permissions:
-        response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
-        response.error_message = ('`grant.permissions` cannot be empty.')
-        return response
-
-      for permission in grant.permissions:
-        logging.info('Granting permission: %r', permission)
+    match request.WhichOneof('verb'):
+      case 'list':
+        what = request.list.WhichOneof('what')
+        response, output = self._execute_command(
+            ['shell', 'pm', 'list', what], timeout=timeout
+        )
+
+        if output:
+          items = output.decode('utf-8').split()
+          # Remove prefix for each item.
+          prefix = {
+              'features': 'feature:',
+              'libraries': 'library:',
+              'packages': 'package:',
+          }[what]
+          items = [x[len(prefix) :] for x in items if x.startswith(prefix)]
+          response.package_manager.list.items.extend(items)
+        response.package_manager.output = output
+      case 'clear':
+        package_name = request.clear.package_name
+        if not package_name:
+          response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
+          response.error_message = (
+              f'Empty PackageManagerRequest.clear.package_name. Got: {request}.'
+          )
+          return response
+
+        args = ['shell', 'pm', 'clear', package_name]
+        if request.clear.user_id:
+          args.insert(3, '-f')
+          args.insert(4, request.clear.user_id)
         response, response.package_manager.output = self._execute_command(
-            ['shell', 'pm', 'grant', grant.package_name, permission],
-            timeout=timeout)
+            args, timeout=timeout
+        )
+      case 'grant':
+        grant = request.grant
+        if not grant.package_name:
+          response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
+          response.error_message = '`grant.package_name` cannot be empty.'
+          return response
+
+        if not grant.permissions:
+          response.status = adb_pb2.AdbResponse.Status.FAILED_PRECONDITION
+          response.error_message = '`grant.permissions` cannot be empty.'
+          return response
+
+        for permission in grant.permissions:
+          logging.info('Granting permission: %r', permission)
+          response, response.package_manager.output = self._execute_command(
+              ['shell', 'pm', 'grant', grant.package_name, permission],
+              timeout=timeout,
+          )
 
     return response
 
-  def _handle_dumpsys(self,
-                      request: adb_pb2.AdbRequest,
-                      timeout: Optional[float] = None) -> adb_pb2.AdbResponse:
+  def _handle_dumpsys(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
     """Handles DumpsysRequest messages.
 
     Args: