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

android_env/components/a11y/a11y_servicer.py

@@ -0,0 +1,199 @@
+# coding=utf-8
+# 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.
+# 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.
+
+"""Accessibility Servicer implementation."""
+
+import asyncio
+from collections.abc import AsyncIterator, Generator, Iterable
+import threading
+
+from absl import logging
+from android_env.proto.a11y import a11y_pb2
+from android_env.proto.a11y import a11y_pb2_grpc
+from android_env.proto.a11y import android_accessibility_forest_pb2
+import grpc
+
+
+class A11yServicer(a11y_pb2_grpc.A11yServiceServicer):
+  """Services the A11yService requests."""
+
+  def __init__(self, latest_forest_only: bool = False):
+    self._received_forests: list[
+        android_accessibility_forest_pb2.AndroidAccessibilityForest
+    ] = []
+    self._received_events: list[a11y_pb2.EventRequest] = []
+    self._lock_forests = threading.Lock()
+    self._lock_events = threading.Lock()
+    self._latest_forest_only = latest_forest_only
+    self._paused = True
+
+    # A11y Forest bookkeeping.
+    self._get_forest = asyncio.Event()  # Whether to request a forest.
+    self._forest_ready = asyncio.Event()  # Whether the forest is ready.
+    self._latest_forest: (
+        android_accessibility_forest_pb2.AndroidAccessibilityForest | None
+    ) = None
+
+  def SendForest(
+      self,
+      request: android_accessibility_forest_pb2.AndroidAccessibilityForest,
+      context: grpc.ServicerContext,
+  ) -> a11y_pb2.ForestResponse:
+    self._process_forest(request)
+    return a11y_pb2.ForestResponse()
+
+  def SendEvent(
+      self,
+      request: a11y_pb2.EventRequest,
+      context: grpc.ServicerContext,
+  ) -> a11y_pb2.EventResponse:
+    self._process_event(request)
+    return a11y_pb2.EventResponse()
+
+  async def Bidi(
+      self,
+      request_iterator: AsyncIterator[a11y_pb2.ClientToServer],
+      context: grpc.aio.ServicerContext,
+  ) -> AsyncIterator[a11y_pb2.ServerToClient]:
+    """Processes incoming ClientToServer requests."""
+
+    logging.info('Starting A11yServicer.Bidi()')
+
+    # Send a dummy message to unblock clients in their loop.
+    yield a11y_pb2.ServerToClient()
+
+    # This block defines two coroutines:
+    #
+    # * `read_client_requests()`
+    # * `check_forest()`
+    #
+    # They cooperate with each other and both populate a queue `q` which is
+    # consumed in a loop below, which actually yields requests which are sent to
+    # the client. The processing finishes when the clients "closes" the
+    # connection, which causes `read_client_requests()` to put a special value,
+    # `STOP_ITERATION`, in the queue.
+
+    # Queue for communicating from coroutines to `Bidi()`.
+    q = asyncio.Queue()
+
+    should_run = True
+
+    async def read_client_requests():
+      """Coroutine for reading client requests."""
+
+      nonlocal should_run
+      async for request in request_iterator:
+        field_name = request.WhichOneof('payload')
+        match field_name:
+          case 'event':
+            self._process_event(request.event)
+          case 'forest':
+            self._latest_forest = request.forest
+            self._forest_ready.set()
+            self._get_forest.clear()  # Reset the `Event`.
+          case _:
+            logging.error('Unknown field %r', field_name)
+        await q.put(a11y_pb2.ServerToClient())
+
+      # Send a special value to stop processing this `Bidi` connection.
+      await q.put('STOP_ITERATION')
+      should_run = False
+
+    async def check_forest():
+      """Coroutine for sending "get forest" requests."""
+
+      nonlocal should_run
+      while should_run:
+        await self._get_forest.wait()
+        await q.put(a11y_pb2.ServerToClient(get_forest={}))
+
+    tasks = asyncio.gather(read_client_requests(), check_forest())
+
+    while should_run:
+      v = await q.get()
+      if v == 'STOP_ITERATION':
+        break
+      else:
+        yield v
+
+    await tasks
+
+    logging.info('Finishing A11yServicer.Bidi()')
+
+  async def get_forest(
+      self,
+  ) -> android_accessibility_forest_pb2.AndroidAccessibilityForest | None:
+    """Issues a request to get the a11y forest from the client."""
+
+    self._get_forest.set()  # Unblocks coroutine to send a request.
+    await self._forest_ready.wait()  # Wait for forest to be ready.
+    self._forest_ready.clear()  # Reset the `Event`.
+    return self._latest_forest
+
+  def gather_forests(
+      self,
+  ) -> list[android_accessibility_forest_pb2.AndroidAccessibilityForest]:
+    forests = []
+    with self._lock_forests:
+      forests = self._received_forests
+      self._received_forests = []
+    return forests
+
+  def gather_events(self) -> list[a11y_pb2.EventRequest]:
+    events = []
+    with self._lock_events:
+      events = self._received_events
+      self._received_events = []
+    return events
+
+  def pause_and_clear(self) -> None:
+    """Temporarily stop receiving events/forests and clear the queue.
+
+    Used when resetting the environment; in this case:
+    - all events/forests that have been received since last timestep are things
+      that happened in the last episode after its `LAST` timestep (so we should
+      ignore them, done by clearing the lists).
+    - we're about to receive a bunch of events/forests just as a result of
+      resetting the environment. We don't want to count these either; thus we
+      temporarily stop receiving new ones.
+    """
+    self._paused = True
+    with self._lock_forests:
+      self._received_forests = []
+    with self._lock_events:
+      self._received_events = []
+
+  def resume(self) -> None:
+    """Start receiving events/forests (e.g., after a reset)."""
+    self._paused = False
+
+  def _process_event(self, event: a11y_pb2.EventRequest) -> None:
+    """Adds the given event to the internal buffer of events."""
+
+    if not self._paused:
+      with self._lock_events:
+        self._received_events.append(event)
+
+  def _process_forest(
+      self, forest: android_accessibility_forest_pb2.AndroidAccessibilityForest
+  ) -> None:
+    """Adds the given forest to the internal buffer of forests."""
+
+    if not self._paused:
+      with self._lock_forests:
+        if self._latest_forest_only:
+          self._received_forests = [forest]
+        else:
+          self._received_forests.append(forest)

android_env/components/a11y/a11y_servicer_test.py

@@ -0,0 +1,224 @@
+# coding=utf-8
+# 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.
+# 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.
+
+"""Tests for a11y_servicer."""
+
+import asyncio
+from collections.abc import AsyncIterator, Iterable
+from typing import TypeVar
+from unittest import IsolatedAsyncioTestCase, mock
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from android_env.components.a11y import a11y_servicer
+from android_env.proto.a11y import a11y_pb2
+from android_env.proto.a11y import android_accessibility_forest_pb2
+import grpc
+
+
+_T = TypeVar('_T')
+
+
+async def _aiter(xs: Iterable[_T]) -> AsyncIterator[_T]:
+  """Utility to make an AsyncIterator from Iterable."""
+
+  for x in xs:
+    yield x
+
+
+def one_window_one_node_forest() -> (
+    android_accessibility_forest_pb2.AndroidAccessibilityForest
+):
+  forest = android_accessibility_forest_pb2.AndroidAccessibilityForest()
+  window = forest.windows.add()
+  node = window.tree.nodes.add()
+  node.class_name = 'foo'
+  node.is_clickable = True
+  node.hint_text = 'Foo hint'
+  return forest
+
+
+def one_window_two_nodes_forest() -> (
+    android_accessibility_forest_pb2.AndroidAccessibilityForest
+):
+  forest = android_accessibility_forest_pb2.AndroidAccessibilityForest()
+  window = forest.windows.add()
+  node = window.tree.nodes.add()
+  node.class_name = 'bar'
+  node.is_clickable = True
+  node.hint_text = 'Bar hint'
+  node = window.tree.nodes.add()
+  node.class_name = 'bar'
+  node.is_clickable = False
+  node.hint_text = 'Bar hint 2'
+  return forest
+
+
+def empty_dict() -> dict[str, str]:
+  return {}
+
+
+def single_item_dict_with_special_chars() -> dict[str, str]:
+  return {'foo': 'bar\r\t\nbaz'}
+
+
+class A11yServicerTest(parameterized.TestCase, IsolatedAsyncioTestCase):
+
+  def test_servicer_sendforest(self):
+    mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)
+    servicer = a11y_servicer.A11yServicer()
+    servicer.resume()
+    response = servicer.SendForest(one_window_one_node_forest(), mock_context)
+    self.assertEqual(response.error, '')
+    response = servicer.SendForest(one_window_two_nodes_forest(), mock_context)
+    self.assertEqual(response.error, '')
+    forests = servicer.gather_forests()
+    self.assertLen(forests, 2)
+    self.assertEqual(forests[0], one_window_one_node_forest())
+    self.assertEqual(forests[1], one_window_two_nodes_forest())
+
+  async def test_servicer_bidi_forests(self):
+    """Checks that the bidirectional interface accepts forests."""
+
+    # Arrange.
+    mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)
+    servicer = a11y_servicer.A11yServicer()
+
+    # Act.
+    servicer.resume()
+    responses = [
+        x
+        async for x in servicer.Bidi(
+            _aiter([
+                a11y_pb2.ClientToServer(
+                    event=a11y_pb2.EventRequest(
+                        event=single_item_dict_with_special_chars()
+                    )
+                ),
+                a11y_pb2.ClientToServer(forest=one_window_two_nodes_forest()),
+            ]),
+            mock_context,
+        )
+    ]
+    forest = await servicer.get_forest()
+
+    # Assert.
+    self.assertEqual(responses[0], a11y_pb2.ServerToClient())
+    self.assertEqual(responses[1], a11y_pb2.ServerToClient())
+    self.assertIsNotNone(forest)
+    self.assertEqual(forest, one_window_two_nodes_forest())
+
+  def test_servicer_sendforest_latest_only(self):
+    mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)
+    servicer = a11y_servicer.A11yServicer(latest_forest_only=True)
+    servicer.resume()
+    response = servicer.SendForest(one_window_one_node_forest(), mock_context)
+    self.assertEqual(response.error, '')
+    response = servicer.SendForest(one_window_two_nodes_forest(), mock_context)
+    self.assertEqual(response.error, '')
+    forests = servicer.gather_forests()
+    self.assertLen(forests, 1)
+    self.assertEqual(forests[0], one_window_two_nodes_forest())
+
+  def test_servicer_sendevent(self):
+    mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)
+    servicer = a11y_servicer.A11yServicer()
+    servicer.resume()
+    response = servicer.SendEvent(
+        a11y_pb2.EventRequest(event=empty_dict()), mock_context
+    )
+    self.assertEqual(response.error, '')
+    response = servicer.SendEvent(
+        a11y_pb2.EventRequest(event=single_item_dict_with_special_chars()),
+        mock_context,
+    )
+    self.assertEqual(response.error, '')
+    events = servicer.gather_events()
+    self.assertLen(events, 2)
+    self.assertEqual(events[0].event, empty_dict())
+    self.assertEqual(events[1].event, single_item_dict_with_special_chars())
+
+  async def test_servicer_bidi_events(self):
+    """Checks that the bidirectional interface accepts events."""
+
+    # Arrange.
+    mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)
+    servicer = a11y_servicer.A11yServicer()
+
+    # Act.
+    servicer.resume()
+    responses = [
+        x
+        async for x in servicer.Bidi(
+            _aiter([
+                a11y_pb2.ClientToServer(
+                    event=a11y_pb2.EventRequest(event=empty_dict())
+                ),
+                a11y_pb2.ClientToServer(
+                    event=a11y_pb2.EventRequest(
+                        event=single_item_dict_with_special_chars()
+                    )
+                ),
+            ]),
+            mock_context,
+        )
+    ]
+    events = servicer.gather_events()
+
+    # Assert.
+    self.assertEqual(responses[0], a11y_pb2.ServerToClient())
+    self.assertEqual(responses[1], a11y_pb2.ServerToClient())
+    self.assertLen(events, 2)
+    self.assertEqual(events[0].event, empty_dict())
+    self.assertEqual(events[1].event, single_item_dict_with_special_chars())
+
+  def test_servicer_pause_and_clear_pauses(self):
+    mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)
+    servicer = a11y_servicer.A11yServicer()
+    servicer.resume()
+    servicer.pause_and_clear()
+    response = servicer.SendEvent(
+        a11y_pb2.EventRequest(event=empty_dict()), mock_context
+    )
+    self.assertEqual(response.error, '')
+    response = servicer.SendForest(one_window_one_node_forest(), mock_context)
+    self.assertEqual(response.error, '')
+    events = servicer.gather_events()
+    self.assertEmpty(events)
+    forests = servicer.gather_forests()
+    self.assertEmpty(forests)
+
+  def test_servicer_pause_and_clear_clears(self):
+    mock_context = mock.create_autospec(grpc.ServicerContext, instance=True)
+    servicer = a11y_servicer.A11yServicer()
+    servicer.resume()
+    response = servicer.SendEvent(
+        a11y_pb2.EventRequest(event=empty_dict()), mock_context
+    )
+    self.assertEqual(response.error, '')
+    response = servicer.SendForest(one_window_one_node_forest(), mock_context)
+    self.assertEqual(
+        response.error,
+        '',
+    )
+    servicer.pause_and_clear()
+    events = servicer.gather_events()
+    self.assertEmpty(events)
+    forests = servicer.gather_forests()
+    self.assertEmpty(forests)
+
+
+if __name__ == '__main__':
+  absltest.main()

android_env/components/action_fns.py

@@ -0,0 +1,132 @@
+# coding=utf-8
+# 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.
+# 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.
+
+"""Functions to convert actions between different components' formats."""
+
+from absl import logging
+from android_env.components import action_type as action_type_lib
+from android_env.components import errors
+from android_env.components import pixel_fns
+from android_env.components.simulators import base_simulator
+import numpy as np
+
+
+def send_action_to_simulator(
+    action: dict[str, np.ndarray],
+    simulator: base_simulator.BaseSimulator,
+    screen_width: int,
+    screen_height: int,
+    num_fingers: int,
+) -> bool:
+  """Sends the selected action to the given simulator.
+
+  The simulator will interpret the action according to `action["action_type"]`.
+  The effect this action triggers in the Android OS will be determined by the
+  currently running application.
+
+  Args:
+    action: action which will get interpreted as a touchscreen event.
+    simulator: The simulator that will receive the action.
+    screen_width: The width of the touchscreen in pixels.
+    screen_height: The height of the touchscreen in pixels.
+    num_fingers: The number of fingers used in this simulator.
+  """
+
+  try:
+    match action['action_type']:
+      # If the action is a TOUCH or LIFT, send a touch event to the simulator.
+      case action_type_lib.ActionType.TOUCH | action_type_lib.ActionType.LIFT:
+        prepared_action = _prepare_touch_action(
+            action, screen_width, screen_height, num_fingers
+        )
+        simulator.send_touch(prepared_action)
+      # If the action is a key event, send a key event to the simulator.
+      case action_type_lib.ActionType.KEYDOWN:
+        simulator.send_key(action['keycode'].item(0), event_type='keydown')
+      case action_type_lib.ActionType.KEYUP:
+        simulator.send_key(action['keycode'].item(0), event_type='keyup')
+      case action_type_lib.ActionType.KEYPRESS:
+        simulator.send_key(action['keycode'].item(0), event_type='keypress')
+  except errors.SendActionError:
+    logging.exception('Unable to execute action: %r', action)
+    return False
+
+  return True
+
+
+def _prepare_touch_action(
+    action: dict[str, np.ndarray],
+    screen_width: int,
+    screen_height: int,
+    num_fingers: int,
+) -> list[tuple[int, int, bool, int]]:
+  """Turns an AndroidEnv action into values that the simulator can interpret.
+
+  Converts float-valued 'touch_position' to integer coordinates corresponding
+  to specific pixels, and 'action_type' to booleans indicating whether the
+  screen is touched at said location or not. The result of this function can
+  be sent directly to the underlying simulator (e.g. the Android Emulator,
+  virtual machine, or a phone).
+
+  Args:
+    action: An action containing 'action_type' and 'touch_position'.
+
+  Returns:
+    A tuple with the format (x: int, y: int, down/up: bool, finger_index: int).
+  """
+
+  touch_events = []
+  for i, finger_action in enumerate(_split_touch_action(action, num_fingers)):
+    is_touch = finger_action['action_type'] == action_type_lib.ActionType.TOUCH
+    touch_position = finger_action['touch_position']
+    touch_pixels = pixel_fns.touch_position_to_pixel_position(
+        touch_position, width_height=(screen_width, screen_height)
+    )
+    touch_events.append((touch_pixels[0], touch_pixels[1], is_touch, i))
+  return touch_events
+
+
+def _split_touch_action(
+    action: dict[str, np.ndarray], num_fingers: int
+) -> list[dict[str, np.ndarray]]:
+  """Splits a multitouch action into a list of single-touch actions."""
+
+  single_touch_actions = [{
+      'action_type': action['action_type'],
+      'touch_position': action['touch_position'],
+  }]
+  for i in range(2, num_fingers + 1):
+    single_touch_actions.append({
+        'action_type': action[f'action_type_{i}'],
+        'touch_position': action[f'touch_position_{i}'],
+    })
+  return single_touch_actions
+
+
+def lift_all_fingers_action(num_fingers: int) -> dict[str, np.ndarray]:
+  """A lift action with each finger."""
+
+  # There's always at least one finger.
+  lift_action = {
+      'action_type': np.array(action_type_lib.ActionType.LIFT),
+      'touch_position': np.array([0, 0]),
+  }
+  # Subsequent fingers have separate dict entries.
+  for i in range(2, num_fingers + 1):
+    lift_action |= {
+        f'action_type_{i}': np.array(action_type_lib.ActionType.LIFT),
+        f'touch_position_{i}': np.array([0, 0]),
+    }
+  return lift_action