diff --git a/android_env/components/coordinator.py b/android_env/components/coordinator.py index 31b9313..c409633 100644 --- a/android_env/components/coordinator.py +++ b/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. diff --git a/android_env/components/simulators/base_simulator.py b/android_env/components/simulators/base_simulator.py index 66b8e79..23f13f7 100644 --- a/android_env/components/simulators/base_simulator.py +++ b/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.""" diff --git a/android_env/components/simulators/emulator/emulator_simulator.py b/android_env/components/simulators/emulator/emulator_simulator.py index 93fee39..8ef551a 100644 --- a/android_env/components/simulators/emulator/emulator_simulator.py +++ b/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.""" diff --git a/android_env/components/specs.py b/android_env/components/specs.py index 9f8a76d..d84f405 100644 --- a/android_env/components/specs.py +++ b/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, diff --git a/android_env/components/task_manager.py b/android_env/components/task_manager.py index c42e5ce..34144ec 100644 --- a/android_env/components/task_manager.py +++ b/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: ( diff --git a/android_env/environment.py b/android_env/environment.py index 9193420..84263ef 100644 --- a/android_env/environment.py +++ b/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 = {}