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

android_env/__init__.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 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.

android_env/components/__init__.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 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.

android_env/components/a11y/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+

android_env/components/a11y/a11y_events.py

@@ -0,0 +1,118 @@
+# 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.
+
+"""Tools for accessing accessibility events."""
+
+from collections.abc import Mapping
+from typing import Any
+
+from absl import logging
+from android_env.proto.a11y import a11y_pb2
+import numpy as np
+
+from google.protobuf import any_pb2
+
+
+_A11Y_EVENT_KEY = 'full_event'
+
+
+def package_events_to_task_extras(
+    events: list[a11y_pb2.EventRequest],
+) -> Mapping[str, np.ndarray]:
+  if not events:
+    return {}
+  events = np.stack(events, axis=0)
+  return {_A11Y_EVENT_KEY: events}
+
+
+def extract_events_from_task_extras(
+    task_extras: Mapping[str, Any] | None = None,
+) -> list[Mapping[str, str]]:
+  """Inspects task_extras and extracts all accessibility events detected.
+
+  Args:
+    task_extras: Task extras forwarded by AndroidEnv. If 'full_event' is not a
+      key in task_extras, then this function returns an empty string. Otherwise,
+      full_event is expected to be list to be a numpy array with one dimension,
+      and contains a list of dictionary describing accessibility events that are
+      present in the given task extras. e.g. 'event_type:
+      TYPE_WINDOW_CONTENT_CHANGED // event_package_name:
+      com.google.android.deskclock // source_class_name:
+      android.widget.ImageView'.
+
+  Returns:
+    List of all events detected
+  """
+  if task_extras is None or _A11Y_EVENT_KEY not in task_extras:
+    return []
+
+  if (
+      not isinstance(task_extras[_A11Y_EVENT_KEY], np.ndarray)
+      or task_extras[_A11Y_EVENT_KEY].ndim != 1
+  ):
+    raise ValueError(
+        f'{_A11Y_EVENT_KEY} task extra should be a numpy array with one'
+        ' dimension.'
+    )
+
+  if task_extras[_A11Y_EVENT_KEY].size == 0:
+    return []
+
+  events = []
+  for e in task_extras[_A11Y_EVENT_KEY]:
+    if isinstance(e, a11y_pb2.EventRequest):
+      events.append(dict(e.event))
+    elif isinstance(e, dict):
+      events.append(e)
+      logging.warning(
+          'The event should come only from the a11y_grpc_wrapper. '
+          'Please verify that the upacking operation has not been '
+          'called twice. See here for full task_extras: %s',
+          task_extras,
+      )
+    elif isinstance(e, any_pb2.Any):
+      ev = a11y_pb2.EventRequest()
+      new_any = any_pb2.Any()
+      new_any.CopyFrom(e)
+      new_any.Unpack(ev)
+      events.append(dict(ev.event))
+
+    else:
+      raise TypeError(
+          f'Unexpected event type: {type(e)}. See here for full '
+          f'task_extras: {task_extras}.'
+      )
+
+  return events
+
+
+def keep_latest_event_only(task_extras: dict[str, Any]):
+  """Removes all a11y events except the last one observed."""
+  if task_extras is None or 'full_event' not in task_extras:
+    return
+
+  if (
+      not isinstance(task_extras[_A11Y_EVENT_KEY], np.ndarray)
+      or task_extras[_A11Y_EVENT_KEY].ndim != 1
+  ):
+    raise ValueError(
+        f'{_A11Y_EVENT_KEY} task extra should be a numpy array with one'
+        ' dimension.'
+    )
+
+  if task_extras[_A11Y_EVENT_KEY].size == 0:
+    return []
+
+  task_extras[_A11Y_EVENT_KEY] = task_extras[_A11Y_EVENT_KEY][-1:]

android_env/components/a11y/a11y_events_test.py

@@ -0,0 +1,173 @@
+# 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_events."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from android_env.components.a11y import a11y_events
+from android_env.proto.a11y import a11y_pb2
+import numpy as np
+
+from google.protobuf import any_pb2
+
+
+def _event_request(d: dict[str, str]) -> a11y_pb2.EventRequest:
+  event_request = a11y_pb2.EventRequest()
+  for k, v in d.items():
+    event_request.event[k] = v
+  return event_request
+
+
+def _event_request_as_any(d: dict[str, str]) -> any_pb2.Any:
+  event_request = _event_request(d)
+  response = any_pb2.Any()
+  response.Pack(event_request)
+  return response
+
+
+class A11yEventsTest(parameterized.TestCase):
+
+  @parameterized.parameters(
+      dict(task_extras={}),
+      dict(
+          task_extras={'no_full_event': [{'1': '1'}, {'2': '2'}, {'3': '3'}]},
+      ),
+      dict(
+          task_extras={'full_event': np.array([])},
+      ),
+      dict(
+          task_extras={},
+      ),
+  )
+  def test_no_events_in_task_extras(self, task_extras):
+    events = a11y_events.extract_events_from_task_extras(task_extras)
+    self.assertEmpty(events)
+
+  @parameterized.parameters(
+      dict(
+          task_extras={'full_event': [{'1': '1'}, {'2': '2'}]},
+          expected_events=[{'1': '1'}, {'2': '2'}],
+      ),
+      dict(
+          task_extras={'full_event': [{}]},
+          expected_events=[{}],
+      ),
+      dict(
+          task_extras={
+              'full_event_wrong_key': [1, 2, 3],
+              'full_event': [{'1': '1'}, {'2': '2'}, {'3': '3'}],
+          },
+          expected_events=[{'1': '1'}, {'2': '2'}, {'3': '3'}],
+      ),
+  )
+  def test_task_extras(self, task_extras, expected_events):
+    event_requests = [_event_request(e) for e in task_extras['full_event']]
+    task_extras['full_event'] = np.stack(event_requests, axis=0)
+    events = a11y_events.extract_events_from_task_extras(task_extras)
+    self.assertEqual(len(events), len(expected_events))
+    for i, event in enumerate(expected_events):
+      self.assertEqual(len(event), len(expected_events[i]))
+      for k, v in event.items():
+        self.assertIn(k, expected_events[i])
+        self.assertEqual(v, expected_events[i][k])
+
+  def test_events_key_has_dict_event_requrests(self):
+    event_requests = [
+        _event_request({'1': '1'}),
+        {'2': '2'},
+        _event_request({'3': '3'}),
+    ]
+    expected_events = [
+        {'1': '1'},
+        {'2': '2'},
+        {'3': '3'},
+    ]
+    task_extras = {'full_event': np.stack(event_requests, axis=0)}
+    events = a11y_events.extract_events_from_task_extras(task_extras)
+    self.assertEqual(len(events), len(expected_events))
+    for i, event in enumerate(expected_events):
+      self.assertEqual(len(event), len(expected_events[i]))
+      for k, v in event.items():
+        self.assertIn(k, expected_events[i])
+        self.assertEqual(v, expected_events[i][k])
+
+  def test_events_key_has__event_requrests_packed_as_any(self):
+    event_requests = [
+        _event_request_as_any({'1': '1'}),
+        {'2': '2'},
+        _event_request_as_any({'3': '3'}),
+    ]
+    expected_events = [
+        {'1': '1'},
+        {'2': '2'},
+        {'3': '3'},
+    ]
+    task_extras = {'full_event': np.stack(event_requests, axis=0)}
+    events = a11y_events.extract_events_from_task_extras(task_extras)
+    self.assertEqual(len(events), len(expected_events))
+    for i, event in enumerate(expected_events):
+      self.assertEqual(len(event), len(expected_events[i]))
+      for k, v in event.items():
+        self.assertIn(k, expected_events[i])
+        self.assertEqual(v, expected_events[i][k])
+
+  def test_events_key_has_non_event_requrests(self):
+    event_requests = [
+        _event_request({'1': '1'}),
+        3,  # Not an even and not a dict.
+        _event_request({'3': '3'}),
+    ]
+    task_extras = {'full_event': np.stack(event_requests, axis=0)}
+    with self.assertRaises(TypeError):
+      _ = a11y_events.extract_events_from_task_extras(task_extras)
+
+  @parameterized.parameters(
+      dict(task_extras={}, expected_extras={}),
+      dict(
+          task_extras={
+              'no_full_event': 42,
+          },
+          expected_extras={
+              'no_full_event': 42,
+          },
+      ),
+      dict(
+          task_extras={'full_event': np.array([1, 2]), 'no_full_event': 43},
+          expected_extras={'full_event': np.array([2]), 'no_full_event': 43},
+      ),
+      dict(
+          task_extras={'full_event': np.array([1, 2, 3])},
+          expected_extras={'full_event': np.array([3])},
+      ),
+      dict(
+          task_extras={'full_event': np.array([]), 'no_full_event': 44},
+          expected_extras={'full_event': np.array([]), 'no_full_event': 44},
+      ),
+  )
+  def test_keep_latest_only(self, task_extras, expected_extras):
+    a11y_events.keep_latest_event_only(task_extras)
+    self.assertEqual(len(task_extras), len(expected_extras))
+    for k, v in task_extras.items():
+      self.assertIn(k, expected_extras)
+      if k == 'full_event':
+        np.testing.assert_array_equal(v, expected_extras['full_event'])
+      else:
+        self.assertEqual(v, expected_extras[k])
+    pass
+
+
+if __name__ == '__main__':
+  absltest.main()

android_env/components/a11y/a11y_forests.py

@@ -0,0 +1,128 @@
+# 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.
+
+"""Tools for accessing accessibility events."""
+
+from collections.abc import Mapping
+from typing import Any
+
+from android_env.proto.a11y import android_accessibility_forest_pb2
+import numpy as np
+
+from google.protobuf import any_pb2
+
+
+_A11Y_FORESTS_KEY = 'accessibility_tree'
+
+
+def package_forests_to_task_extras(
+    forests: list[android_accessibility_forest_pb2.AndroidAccessibilityForest],
+) -> Mapping[str, np.ndarray]:
+  if not forests:
+    return {}
+  forests = np.stack(forests, axis=0)
+  return {_A11Y_FORESTS_KEY: forests}
+
+
+def task_extras_has_forests(task_extras: Mapping[str, Any]) -> bool:
+  """Checks that the task_extras has any a11y forest information."""
+  if _A11Y_FORESTS_KEY not in task_extras:
+    return False
+
+  payload = task_extras[_A11Y_FORESTS_KEY]
+  if not isinstance(payload, np.ndarray) or payload.ndim != 1:
+    raise ValueError(
+        f'{_A11Y_FORESTS_KEY} task extra should be a numpy array with one'
+        f' dimension. payload: {payload}'
+    )
+
+  if payload.size == 0:
+    return False
+
+  if any(isinstance(f, any_pb2.Any) for f in payload):
+    # Forests were packed as Any.
+    return True
+
+  return any(
+      isinstance(f, android_accessibility_forest_pb2.AndroidAccessibilityForest)
+      for f in payload
+  )
+
+
+def convert_to_forest(
+    forest: android_accessibility_forest_pb2.AndroidAccessibilityForest
+    | any_pb2.Any
+    | None,
+) -> android_accessibility_forest_pb2.AndroidAccessibilityForest | None:
+  """Takes an object and attempts to convert it to a forest."""
+  if forest is None:
+    return None
+
+  if isinstance(forest, any_pb2.Any):
+    output = android_accessibility_forest_pb2.AndroidAccessibilityForest()
+    new_any = any_pb2.Any()
+    new_any.CopyFrom(forest)
+    new_any.Unpack(output)
+    return output
+  elif isinstance(
+      forest, android_accessibility_forest_pb2.AndroidAccessibilityForest
+  ):
+    return forest
+  else:
+    return None
+
+
+def extract_forests_from_task_extras(
+    task_extras: Mapping[str, Any] | None = None,
+) -> list[android_accessibility_forest_pb2.AndroidAccessibilityForest]:
+  """Inspects task_extras and extracts all accessibility forests detected.
+
+  Args:
+    task_extras: Task extras forwarded by AndroidEnv. If 'full_event' is not a
+      key in task_extras, then this function returns an empty string. Otherwise,
+      full_event is expected to be list to be a numpy array with one dimension,
+      and contains a list of dictionary describing accessibility forests that
+      are present in the given task extras.
+
+  Returns:
+    List of all forests detected
+  """
+  if task_extras is None or not task_extras_has_forests(task_extras):
+    return []
+
+  forests = []
+  for f in task_extras[_A11Y_FORESTS_KEY]:
+    f = convert_to_forest(f)
+    if f is not None:
+      forests.append(f)
+  return forests
+
+
+def keep_latest_forest_only(task_extras: dict[str, Any]):
+  """Removes all a11y forests except the last one observed."""
+  if _A11Y_FORESTS_KEY not in task_extras.keys():
+    return
+
+  payload = task_extras[_A11Y_FORESTS_KEY]
+  if not isinstance(payload, np.ndarray) or payload.ndim != 1:
+    raise ValueError(
+        f'{_A11Y_FORESTS_KEY} task extra should be a numpy array with one'
+        f' dimension. payload: {payload}'
+    )
+
+  if payload.size == 0:
+    return
+
+  task_extras[_A11Y_FORESTS_KEY] = payload[-1:]

android_env/components/a11y/a11y_forests_test.py

@@ -0,0 +1,237 @@
+# 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_forests."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from android_env.components.a11y import a11y_forests
+from android_env.proto.a11y import android_accessibility_forest_pb2
+import numpy as np
+
+from google.protobuf import any_pb2
+
+
+def _pack_any(proto_message) -> any_pb2.Any:
+  response = any_pb2.Any()
+  response.Pack(proto_message)
+  return response
+
+
+def _empty_forest() -> (
+    android_accessibility_forest_pb2.AndroidAccessibilityForest
+):
+  return android_accessibility_forest_pb2.AndroidAccessibilityForest()
+
+
+def _one_empty_window_forest() -> (
+    android_accessibility_forest_pb2.AndroidAccessibilityForest
+):
+  forest = android_accessibility_forest_pb2.AndroidAccessibilityForest()
+  forest.windows.add()
+  return forest
+
+
+def _two_window_forest() -> (
+    android_accessibility_forest_pb2.AndroidAccessibilityForest
+):
+  forest = android_accessibility_forest_pb2.AndroidAccessibilityForest()
+  window = forest.windows.add()
+  window.tree.nodes.add(
+      class_name='foo', is_clickable=True, hint_text='Foo hint'
+  )
+  forest.windows.add()
+  return forest
+
+
+class A11YForestsTest(parameterized.TestCase):
+
+  @parameterized.parameters(
+      dict(task_extras={}, expected_forests=[], convert_to_np=[]),
+      dict(
+          task_extras={'accessibility_tree': []},
+          convert_to_np=['accessibility_tree'],
+          expected_forests=[],
+      ),
+      dict(
+          task_extras={
+              'not_accessibility_tree': [
+                  _empty_forest(),
+                  _one_empty_window_forest(),
+                  _two_window_forest(),
+              ],
+          },
+          convert_to_np=['not_accessibility_tree'],
+          expected_forests=[],
+      ),
+      dict(
+          task_extras={
+              'accessibility_tree': [
+                  _empty_forest(),
+                  {'not_a_forest_key': 'nor_a_forest_value'},
+                  _two_window_forest(),
+              ]
+          },
+          convert_to_np=['accessibility_tree'],
+          expected_forests=[_empty_forest(), _two_window_forest()],
+      ),
+      dict(
+          task_extras={
+              'accessibility_tree': [
+                  {'not_a_forest_key': 'nor_a_forest_value'},
+                  3,
+                  4,
+                  {'not_a_forest_key': _empty_forest()},
+              ],
+          },
+          convert_to_np=['accessibility_tree'],
+          expected_forests=[],
+      ),
+      dict(
+          task_extras={'accessibility_tree': []},
+          convert_to_np=['accessibility_tree'],
+          expected_forests=[],
+      ),
+      dict(
+          task_extras={
+              'accessibility_tree_wrong_key': [1, 2, 3],
+              'accessibility_tree': [
+                  _empty_forest(),
+                  None,
+                  None,
+                  _one_empty_window_forest(),
+                  _two_window_forest(),
+              ],
+          },
+          convert_to_np=['accessibility_tree', 'accessibility_tree_wrong_key'],
+          expected_forests=[
+              _empty_forest(),
+              _one_empty_window_forest(),
+              _two_window_forest(),
+          ],
+      ),
+      dict(
+          task_extras={
+              'accessibility_tree_wrong_key': [1, 2, 3],
+              'accessibility_tree': [
+                  None,
+                  _pack_any(_empty_forest()),
+                  _pack_any(_one_empty_window_forest()),
+                  _pack_any(_two_window_forest()),
+              ],
+          },
+          convert_to_np=['accessibility_tree', 'accessibility_tree_wrong_key'],
+          expected_forests=[
+              _empty_forest(),
+              _one_empty_window_forest(),
+              _two_window_forest(),
+          ],
+      ),
+      dict(
+          task_extras={
+              'accessibility_tree': [
+                  _pack_any(_empty_forest()),
+                  {'not_a_forest_key': 'nor_a_forest_value'},
+                  None,
+                  _two_window_forest(),
+                  None,
+              ]
+          },
+          convert_to_np=['accessibility_tree'],
+          expected_forests=[_empty_forest(), _two_window_forest()],
+      ),
+  )
+  def test_task_extras(self, task_extras, expected_forests, convert_to_np):
+    for k in convert_to_np:
+      if task_extras[k]:
+        task_extras[k] = np.stack(task_extras[k], axis=0)
+      else:
+        task_extras[k] = np.array([])
+    forests = a11y_forests.extract_forests_from_task_extras(task_extras)
+    self.assertEqual(len(forests), len(expected_forests))
+    for idx, f in enumerate(forests):
+      self.assertEqual(f, expected_forests[idx])
+
+  @parameterized.parameters(
+      dict(task_extras={}, expected_extras={}),
+      dict(
+          task_extras={
+              'no_accessibility_tree': 42,
+          },
+          expected_extras={
+              'no_accessibility_tree': 42,
+          },
+      ),
+      dict(
+          task_extras={'accessibility_tree': []},
+          expected_extras={'accessibility_tree': []},
+      ),
+      dict(
+          task_extras={
+              'accessibility_tree': [
+                  _empty_forest(),
+                  _one_empty_window_forest(),
+              ],
+              'no_accessibility_tree': 43,
+          },
+          expected_extras={
+              'accessibility_tree': [_one_empty_window_forest()],
+              'no_accessibility_tree': 43,
+          },
+      ),
+      dict(
+          task_extras={
+              'accessibility_tree': [
+                  _empty_forest(),
+                  _one_empty_window_forest(),
+                  _two_window_forest(),
+              ]
+          },
+          expected_extras={'accessibility_tree': [_two_window_forest()]},
+      ),
+      dict(
+          task_extras={
+              'accessibility_tree': [],
+              'no_accessibility_tree': 44,
+          },
+          expected_extras={
+              'accessibility_tree': [],
+              'no_accessibility_tree': 44,
+          },
+      ),
+  )
+  def test_keep_latest_only(self, task_extras, expected_extras):
+    if 'accessibility_tree' in task_extras:
+      if task_extras['accessibility_tree']:
+        task_extras['accessibility_tree'] = np.stack(
+            task_extras['accessibility_tree'], axis=0
+        )
+      else:
+        task_extras['accessibility_tree'] = np.array([])
+
+    a11y_forests.keep_latest_forest_only(task_extras)
+    self.assertSameElements(task_extras.keys(), expected_extras.keys())
+    for k in task_extras.keys():
+      if k == 'accessibility_tree':
+        self.assertEqual(len(task_extras[k]), len(expected_extras[k]))
+        for idx, f in enumerate(task_extras[k]):
+          self.assertEqual(f, expected_extras[k][idx])
+      else:
+        self.assertEqual(task_extras[k], expected_extras[k])
+    pass
+
+
+if __name__ == '__main__':
+  absltest.main()