Adopt PEP 591 typing upgrades and readability improvements in AndroidEnv.

android_env/components/coordinator.py

@@ -15,9 +15,11 @@
 
 """Coordinator handles interaction between internal components of AndroidEnv."""
 
+from __future__ import annotations
+
 import copy
 import time
-from typing import Any
+from typing import Any, Final, final
 
 from absl import logging
 from android_env.components import action_fns
@@ -33,6 +35,7 @@
 import numpy as np
 
 
+@final
 class Coordinator:
   """Handles interaction between internal components of AndroidEnv."""
 
@@ -50,10 +53,14 @@ def __init__(
       task_manager: The TaskManager, responsible for coordinating RL tasks.
       config: Settings to customize this Coordinator.
     """
-    self._simulator = simulator
-    self._task_manager = task_manager
-    self._config = config or config_classes.CoordinatorConfig()
-    self._device_settings = device_settings
+    self._simulator: Final[base_simulator.BaseSimulator] = simulator
+    self._task_manager: Final[task_manager_lib.TaskManager] = task_manager
+    self._config: Final[config_classes.CoordinatorConfig] = (
+        config or config_classes.CoordinatorConfig()
+    )
+    self._device_settings: Final[device_settings_lib.DeviceSettings] = (
+        device_settings
+    )
     self._adb_call_parser: adb_call_parser.AdbCallParser | None = None
 
     # Initialize stats.

android_env/components/simulators/base_simulator.py

@@ -15,10 +15,13 @@
 
 """A base class for talking to different types of Android simulators."""
 
+from __future__ import annotations
+
 import abc
 from collections.abc import Callable
 import threading
 import time
+from typing import final
 
 from absl import logging
 from android_env.components import adb_controller
@@ -65,6 +68,7 @@ def create_adb_controller(self) -> adb_controller.AdbController:
   def create_log_stream(self) -> log_stream.LogStream:
     """Creates a stream of logs from the simulator."""
 
+  @final
   def launch(self) -> None:
     """Starts the simulator."""
 
@@ -148,6 +152,7 @@ def save_state(
     """
     raise NotImplementedError('This simulator does not support save_state()')
 
+  @final
   def get_screenshot(self) -> np.ndarray:
     """Returns pixels representing the current screenshot of the simulator."""
 
@@ -174,6 +179,7 @@ def close(self):
       self._interaction_thread.join()
 
 
+@final
 class InteractionThread(threading.Thread):
   """A thread that gets screenshot in the background."""
 

android_env/components/simulators/emulator/emulator_simulator.py

@@ -15,12 +15,14 @@
 
 """A class that manages an Android Emulator."""
 
+from __future__ import annotations
+
 from collections.abc import Callable, Mapping
 import functools
 import os
 import platform
 import time
-from typing import Any
+from typing import Any, Final, final
 
 from absl import logging
 from android_env.components import adb_controller
@@ -42,10 +44,10 @@
 from google.protobuf import empty_pb2
 
 
-_DEFAULT_SNAPSHOT_NAME: str = 'default_snapshot'
+_DEFAULT_SNAPSHOT_NAME: Final[str] = 'default_snapshot'
 
-_KEYCODE_TYPE_BY_SYSTEM: Mapping[
-    str, emulator_controller_pb2.KeyboardEvent.KeyCodeType
+_KEYCODE_TYPE_BY_SYSTEM: Final[
+    Mapping[str, emulator_controller_pb2.KeyboardEvent.KeyCodeType]
 ] = immutabledict.immutabledict({
     'Linux': emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,
     'Windows': emulator_controller_pb2.KeyboardEvent.KeyCodeType.Win,
@@ -89,6 +91,7 @@ def wrapper(self: 'EmulatorSimulator', *args: Any, **kwargs: Any) -> Any:
   return wrapper
 
 
+@final
 class EmulatorSimulator(base_simulator.BaseSimulator):
   """Controls an Android Emulator."""
 

android_env/components/specs.py

@@ -15,13 +15,19 @@
 
 """Base specs for AndroidEnv."""
 
+from __future__ import annotations
+
+from typing import Any, Final, Mapping
+
 from android_env.components import action_type
 from android_env.proto import task_pb2
 from dm_env import specs
 import numpy as np
 
 
-_PROTO_DTYPE_TO_NUMPY_DTYPE = {
+_PROTO_DTYPE_TO_NUMPY_DTYPE: Final[
+    Mapping[task_pb2.ArraySpec.DataType.ValueType, np.dtype | type[Any]]
+] = {
     task_pb2.ArraySpec.DataType.FLOAT: np.float32,
     task_pb2.ArraySpec.DataType.DOUBLE: np.float64,
     task_pb2.ArraySpec.DataType.INT8: np.int8,

android_env/components/task_manager.py

@@ -15,6 +15,8 @@
 
 """TaskManager handles all events and information related to the task."""
 
+from __future__ import annotations
+
 import ast
 import collections
 from collections.abc import Callable, Iterable
@@ -25,7 +27,7 @@
 import re
 import threading
 import time
-from typing import Any
+from typing import Any, Final, final
 
 from absl import logging
 from android_env.components import adb_call_parser as adb_call_parser_lib
@@ -40,6 +42,7 @@
 import numpy as np
 
 
+@final
 class TaskManager:
   """Handles all events and information related to the task."""
 
@@ -55,9 +58,11 @@ def __init__(
       config: Configuration for this instance.
     """
 
-    self._task = task
-    self._config = config or config_classes.TaskManagerConfig()
-    self._lock = threading.Lock()
+    self._task: Final[task_pb2.Task] = task
+    self._config: Final[config_classes.TaskManagerConfig] = (
+        config or config_classes.TaskManagerConfig()
+    )
+    self._lock: Final[threading.Lock] = threading.Lock()
     self._logcat_thread: logcat_thread.LogcatThread | None = None
     self._dumpsys_thread: dumpsys_thread.DumpsysThread | None = None
     self._setup_step_interpreter: (

android_env/environment.py

@@ -15,7 +15,9 @@
 
 """Android environment implementation."""
 
-from typing import Any
+from __future__ import annotations
+
+from typing import Any, Final, final
 
 from absl import logging
 from android_env import env_interface
@@ -29,6 +31,7 @@
 import numpy as np
 
 
+@final
 class AndroidEnv(env_interface.AndroidEnvInterface):
   """An RL environment that interacts with Android apps."""
 
@@ -40,9 +43,9 @@ def __init__(
   ):
     """Initializes the state of this AndroidEnv object."""
 
-    self._simulator = simulator
-    self._coordinator = coordinator
-    self._task_manager = task_manager
+    self._simulator: Final[base_simulator.BaseSimulator] = simulator
+    self._coordinator: Final[coordinator_lib.Coordinator] = coordinator
+    self._task_manager: Final[task_manager_lib.TaskManager] = task_manager
     self._latest_action = {}
     self._latest_observation = {}
     self._latest_extras = {}