Small typing improvements and fixes

Small miscellaneous Python typing improvements and fixes that wouldn't
fit into any other change.

Bug: b/285950835
Change-Id: Idf957cfcdd268a79e0ea8633aec9d3e3244d29f7
Reviewed-on: https://fuchsia-review.googlesource.com/c/antlion/+/885955
Reviewed-by: Patrick Lu <patricklu@google.com>
Commit-Queue: Auto-Submit <auto-submit@fuchsia-infra.iam.gserviceaccount.com>
Fuchsia-Auto-Submit: Sam Balana <sbalana@google.com>

packages/antlion/controllers/iperf_server.py

diff --git a/packages/antlion/controllers/iperf_server.py b/packages/antlion/controllers/iperf_server.py
index 349792b..ee552c5 100755
--- a/packages/antlion/controllers/iperf_server.py
+++ b/packages/antlion/controllers/iperf_server.py
@@ -22,6 +22,7 @@
 import subprocess
 import threading
 import time
+from typing import Union
 
 from antlion import context
 from antlion import logger as acts_logger
@@ -131,7 +132,9 @@
             "sum_received" in self.result["end"] or "sum" in self.result["end"]
         )
 
-    def _get_reporting_speed(self, network_speed_in_bits_per_second):
+    def _get_reporting_speed(
+        self, network_speed_in_bits_per_second: Union[int, float]
+    ) -> float:
         """Sets the units for the network speed reporting based on how the
         object was initiated.  Defaults to Megabytes per second.  Currently
         supported, bits per second (bits), kilobits per second (kbits), megabits

packages/antlion/event/event_bus.py

diff --git a/packages/antlion/event/event_bus.py b/packages/antlion/event/event_bus.py
index 1b210ca..c9ec9f0 100644
--- a/packages/antlion/event/event_bus.py
+++ b/packages/antlion/event/event_bus.py
@@ -292,5 +292,5 @@
             self.event_type, self.func, filter_fn=self.filter_fn, order=self.order
         )
 
-    def __exit__(self, *unused):
+    def __exit__(self, *_):
         _event_bus.unregister(self.registration_id)

packages/antlion/test_utils/wifi/base_test.py

diff --git a/packages/antlion/test_utils/wifi/base_test.py b/packages/antlion/test_utils/wifi/base_test.py
index 03ddb0b..c0de3e84 100644
--- a/packages/antlion/test_utils/wifi/base_test.py
+++ b/packages/antlion/test_utils/wifi/base_test.py
@@ -22,6 +22,7 @@
 from typing import Any, Dict, List, Optional, TypedDict, TypeVar
 
 from mobly.base_test import STAGE_NAME_TEARDOWN_CLASS
+from mobly.config_parser import TestRunConfig
 from mobly.records import TestResultRecord
 
 from antlion import context, controllers, utils
@@ -71,7 +72,7 @@
 
 
 class WifiBaseTest(AntlionBaseTest):
-    def __init__(self, configs: List[Any]) -> None:
+    def __init__(self, configs: TestRunConfig) -> None:
         super().__init__(configs)
         self.enable_packet_log = False
         self.packet_log_2g = hostapd_constants.AP_DEFAULT_CHANNEL_2G

packages/antlion/utils.py

diff --git a/packages/antlion/utils.py b/packages/antlion/utils.py
index a8e78d9..9c0c4ae 100755
--- a/packages/antlion/utils.py
+++ b/packages/antlion/utils.py
@@ -35,7 +35,7 @@
 import traceback
 import zipfile
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any
+from typing import TYPE_CHECKING, Any, Dict, List, Union
 
 import mobly.keys as mobly_keys
 import yaml
@@ -44,6 +44,11 @@
 from antlion.controllers.adb_lib.error import AdbError
 from antlion.libs.proc import job
 
+if TYPE_CHECKING:
+    from antlion.controllers.android_device import AndroidDevice
+    from antlion.controllers.fuchsia_device import FuchsiaDevice
+    from antlion.controllers.utils_lib.ssh.connection import SshConnection
+
 # File name length is limited to 255 chars on some OS, so we need to make sure
 # the file names we output fits within the limit.
 MAX_FILENAME_LEN = 255
@@ -1487,7 +1492,10 @@
     return str(uc_string).encode("ASCII")
 
 
-def get_interface_ip_addresses(comm_channel, interface):
+def get_interface_ip_addresses(
+    comm_channel: Union["AndroidDevice", "SshConnection", "FuchsiaDevice"],
+    interface: str,
+) -> Dict[str, List[str]]:
     """Gets all of the ip addresses, ipv4 and ipv6, associated with a
        particular interface name.
 

pyproject.toml

diff --git a/pyproject.toml b/pyproject.toml
index 391c35d..0354c05 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -127,6 +127,8 @@
     "packages/antlion/test_utils/wifi/wifi_test_utils.py",
     "packages/antlion/utils.py",
     "setup.py",
+
+    "stubs/mobly/",
 ]
 
 [tool.vulture]

tests/wlan/performance/ChannelSweepTest.py

diff --git a/tests/wlan/performance/ChannelSweepTest.py b/tests/wlan/performance/ChannelSweepTest.py
index 5173c11..cbf641a 100644
--- a/tests/wlan/performance/ChannelSweepTest.py
+++ b/tests/wlan/performance/ChannelSweepTest.py
@@ -20,9 +20,10 @@
 from dataclasses import dataclass
 from pathlib import Path
 from statistics import pstdev
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 from mobly import asserts, test_runner
+from mobly.config_parser import TestRunConfig
 
 from antlion import utils
 from antlion.controllers.access_point import setup_ap
@@ -108,7 +109,7 @@
     Note: Performance tests should be done in isolated testbed.
     """
 
-    def __init__(self, configs: List[Any]) -> None:
+    def __init__(self, configs: TestRunConfig) -> None:
         super().__init__(configs)
         self.channel_throughput: ChannelThroughputMap = {}