Add type hints to hardware functions

scripts/wk/hw/audio.py

@@ -12,13 +12,13 @@ LOG = logging.getLogger(__name__)
 
 
 # Functions
-def audio_test():
+def audio_test() -> None:
   """Run an OS-specific audio test."""
   if PLATFORM == 'Linux':
     audio_test_linux()
 
 
-def audio_test_linux():
+def audio_test_linux() -> None:
   """Run an audio test using amixer and speaker-test."""
   LOG.info('Audio Test')
 

scripts/wk/hw/benchmark.py

@@ -38,7 +38,7 @@ class DeviceTooSmallError(RuntimeError):
 
 
 # Functions
-def calc_io_dd_values(dev_size):
+def calc_io_dd_values(dev_size) -> dict[str, int]:
   """Calculate I/O benchmark dd values, returns dict.
 
   Calculations:
@@ -90,8 +90,8 @@ def calc_io_dd_values(dev_size):
     }
 
 
-def check_io_results(test_obj, rate_list, graph_width):
+def check_io_results(test_obj, rate_list, graph_width) -> None:
-  """Generate colored report using rate_list, returns list of str."""
+  """Check I/O restuls and generate report using rate_list."""
   avg_read = sum(rate_list) / len(rate_list)
   min_read = min(rate_list)
   max_read = max(rate_list)
@@ -135,7 +135,7 @@ def check_io_results(test_obj, rate_list, graph_width):
     test_obj.set_status('Unknown')
 
 
-def run_io_test(test_obj, log_path):
+def run_io_test(test_obj, log_path) -> None:
   """Run I/O benchmark and handle exceptions."""
   dev_path = test_obj.dev.path
   if PLATFORM == 'Darwin':

scripts/wk/hw/cpu.py

@@ -5,6 +5,8 @@ import logging
 import re
 import subprocess
 
+from typing import TextIO
+
 from wk import exe
 from wk.cfg.hw import CPU_FAILURE_TEMP
 from wk.os.mac import set_fans as macos_set_fans
@@ -19,10 +21,11 @@ from wk.tmux import respawn_pane as tmux_respawn_pane
 
 # STATIC VARIABLES
 LOG = logging.getLogger(__name__)
+SysbenchType = tuple[subprocess.Popen, TextIO]
 
 
 # Functions
-def check_cooling_results(test_obj, sensors, run_sysbench=False):
+def check_cooling_results(test_obj, sensors, run_sysbench=False) -> None:
   """Check cooling results and update test_obj."""
   max_temp = sensors.cpu_max_temp()
   temp_labels = ['Idle', 'Max', 'Cooldown']
@@ -44,12 +47,12 @@ def check_cooling_results(test_obj, sensors, run_sysbench=False):
     test_obj.report.append(f'  {line}')
 
 
-def check_mprime_results(test_obj, working_dir):
+def check_mprime_results(test_obj, working_dir) -> None:
   """Check mprime log files and update test_obj."""
   passing_lines = {}
   warning_lines = {}
 
-  def _read_file(log_name):
+  def _read_file(log_name) -> list[str]:
     """Read file and split into lines, returns list."""
     lines = []
     try:
@@ -99,7 +102,7 @@ def check_mprime_results(test_obj, working_dir):
     test_obj.report.append(color_string('  Unknown result', 'YELLOW'))
 
 
-def start_mprime(working_dir, log_path):
+def start_mprime(working_dir, log_path) -> subprocess.Popen:
   """Start mprime and save filtered output to log, returns Popen object."""
   set_apple_fan_speed('max')
   proc_mprime = subprocess.Popen( # pylint: disable=consider-using-with
@@ -124,7 +127,7 @@ def start_mprime(working_dir, log_path):
   return proc_mprime
 
 
-def start_sysbench(sensors, sensors_out, log_path, pane):
+def start_sysbench(sensors, sensors_out, log_path, pane) -> SysbenchType:
   """Start sysbench, returns tuple with Popen object and file handle."""
   set_apple_fan_speed('max')
   sysbench_cmd = [
@@ -156,7 +159,7 @@ def start_sysbench(sensors, sensors_out, log_path, pane):
   return (proc_sysbench, filehandle_sysbench)
 
 
-def set_apple_fan_speed(speed):
+def set_apple_fan_speed(speed) -> None:
   """Set Apple fan speed."""
   cmd = None
 
@@ -179,7 +182,7 @@ def set_apple_fan_speed(speed):
     exe.run_program(cmd, check=False)
 
 
-def stop_mprime(proc_mprime):
+def stop_mprime(proc_mprime) -> None:
   """Stop mprime gracefully, then forcefully as needed."""
   proc_mprime.terminate()
   try:
@@ -189,7 +192,7 @@ def stop_mprime(proc_mprime):
   set_apple_fan_speed('auto')
 
 
-def stop_sysbench(proc_sysbench, filehandle_sysbench):
+def stop_sysbench(proc_sysbench, filehandle_sysbench) -> None:
   """Stop sysbench."""
   proc_sysbench.terminate()
   try:

scripts/wk/hw/diags.py

@@ -100,7 +100,7 @@ class State():
     self.init_tmux()
     exe.start_thread(self.fix_tmux_layout_loop)
 
-  def abort_testing(self):
+  def abort_testing(self) -> None:
     """Set unfinished tests as aborted and cleanup tmux panes."""
     for group in self.test_groups:
       for test in group.test_objects:
@@ -119,7 +119,7 @@ class State():
           tmux.kill_pane(_id)
         self.panes.pop(key)
 
-  def disk_safety_checks(self, prep=False, wait_for_self_tests=True):
+  def disk_safety_checks(self, prep=False, wait_for_self_tests=True) -> None:
     # pylint: disable=too-many-branches,too-many-statements
     """Run disk safety checks."""
     self_tests_in_progress = False
@@ -204,8 +204,7 @@ class State():
       std.sleep(60)
       self.disk_safety_checks(wait_for_self_tests=False)
 
-  def fix_tmux_layout(self, forced=True):
+  def fix_tmux_layout(self, forced=True) -> None:
-    # pylint: disable=unused-argument
     """Fix tmux layout based on cfg.hw.TMUX_LAYOUT."""
     try:
       tmux.fix_layout(self.panes, self.layout, forced=forced)
@@ -213,7 +212,7 @@ class State():
       # Assuming self.panes changed while running
       pass
 
-  def fix_tmux_layout_loop(self):
+  def fix_tmux_layout_loop(self) -> None:
     """Fix tmux layout on a loop.
 
     NOTE: This should be called as a thread.
@@ -222,7 +221,7 @@ class State():
       self.fix_tmux_layout(forced=False)
       std.sleep(1)
 
-  def init_diags(self, menu):
+  def init_diags(self, menu) -> None:
     """Initialize diagnostic pass."""
 
     # Reset objects
@@ -293,7 +292,7 @@ class State():
     # Run safety checks
     self.disk_safety_checks(prep=True)
 
-  def init_tmux(self):
+  def init_tmux(self) -> None:
     """Initialize tmux layout."""
     tmux.kill_all_panes()
 
@@ -322,7 +321,7 @@ class State():
       text=' ',
       )
 
-  def save_debug_reports(self):
+  def save_debug_reports(self) -> None:
     """Save debug reports to disk."""
     LOG.info('Saving debug reports')
     debug_dir = pathlib.Path(f'{self.log_dir}/debug')
@@ -368,7 +367,7 @@ class State():
         _f.write(f'\n{test.name}:\n')
         _f.write('\n'.join(debug.generate_object_report(test, indent=1)))
 
-  def update_clock(self):
+  def update_clock(self) -> None:
     """Update 'Started' pane following clock sync."""
     tmux.respawn_pane(
       pane_id=self.panes['Started'],
@@ -379,7 +378,7 @@ class State():
         ),
       )
 
-  def update_progress_pane(self):
+  def update_progress_pane(self) -> None:
     """Update progress pane."""
     report = []
     width = cfg.hw.TMUX_SIDE_WIDTH
@@ -401,13 +400,13 @@ class State():
     with open(out_path, 'w', encoding='utf-8') as _f:
       _f.write('\n'.join(report))
 
-  def update_top_pane(self, text):
+  def update_top_pane(self, text) -> None:
     """Update top pane with text."""
     tmux.respawn_pane(self.panes['Top'], text=f'{self.top_text}\n{text}')
 
 
 # Functions
-def build_menu(cli_mode=False, quick_mode=False):
+def build_menu(cli_mode=False, quick_mode=False) -> std.Menu:
   # pylint: disable=too-many-branches
   """Build main menu, returns wk.std.Menu."""
   menu = std.Menu(title=None)
@@ -459,7 +458,7 @@ def build_menu(cli_mode=False, quick_mode=False):
   return menu
 
 
-def cpu_stress_tests(state, test_objects):
+def cpu_stress_tests(state, test_objects) -> None:
   # pylint: disable=too-many-statements
   """CPU & cooling check using Prime95 and Sysbench."""
   LOG.info('CPU Test (Prime95)')
@@ -587,7 +586,7 @@ def cpu_stress_tests(state, test_objects):
     raise std.GenericAbort('Aborted')
 
 
-def disk_attribute_check(state, test_objects):
+def disk_attribute_check(state, test_objects) -> None:
   """Disk attribute check."""
   LOG.info('Disk Attribute Check')
   for test in test_objects:
@@ -607,7 +606,7 @@ def disk_attribute_check(state, test_objects):
   state.update_progress_pane()
 
 
-def disk_io_benchmark(state, test_objects, skip_usb=True):
+def disk_io_benchmark(state, test_objects, skip_usb=True) -> None:
   # pylint: disable=too-many-statements
   """Disk I/O benchmark using dd."""
   LOG.info('Disk I/O Benchmark (dd)')
@@ -671,7 +670,7 @@ def disk_io_benchmark(state, test_objects, skip_usb=True):
     raise std.GenericAbort('Aborted')
 
 
-def disk_self_test(state, test_objects):
+def disk_self_test(state, test_objects) -> None:
   # pylint: disable=too-many-statements
   """Disk self-test if available."""
   LOG.info('Disk Self-Test(s)')
@@ -729,8 +728,8 @@ def disk_self_test(state, test_objects):
     raise std.GenericAbort('Aborted')
 
 
-def disk_surface_scan(state, test_objects):
+def disk_surface_scan(state, test_objects) -> None:
-  # pylint: disable=too-many-branches,too-many-statements
+  # pylint: disable=too-many-branches
   """Read-only disk surface scan using badblocks."""
   LOG.info('Disk Surface Scan (badblocks)')
   aborted = False
@@ -810,7 +809,7 @@ def disk_surface_scan(state, test_objects):
     raise std.GenericAbort('Aborted')
 
 
-def main():
+def main() -> None:
   # pylint: disable=too-many-branches
   """Main function for hardware diagnostics."""
   args = docopt(DOCSTRING)
@@ -883,7 +882,7 @@ def main():
     state.update_top_pane('Main Menu')
 
 
-def print_countdown(proc, seconds):
+def print_countdown(proc, seconds) -> None:
   """Print countdown to screen while proc is alive."""
   for i in range(seconds):
     sec_left = (seconds - i) % 60
@@ -910,7 +909,7 @@ def print_countdown(proc, seconds):
   print('')
 
 
-def run_diags(state, menu, quick_mode=False):
+def run_diags(state, menu, quick_mode=False) -> None:
   """Run selected diagnostics."""
   aborted = False
   atexit.register(state.save_debug_reports)
@@ -964,7 +963,7 @@ def run_diags(state, menu, quick_mode=False):
     std.pause('Press Enter to return to main menu...')
 
 
-def show_results(state):
+def show_results(state) -> None:
   """Show test results by device."""
   std.sleep(0.5)
   std.clear_screen()
@@ -993,7 +992,7 @@ def show_results(state):
       std.print_standard(' ')
 
 
-def sync_clock():
+def sync_clock() -> None:
   """Sync clock under macOS using sntp."""
   cmd = ['sudo', 'sntp', '-Ss', 'us.pool.ntp.org']
   proc = exe.run_program(cmd, check=False)

scripts/wk/hw/keyboard.py

@@ -12,7 +12,7 @@ LOG = logging.getLogger(__name__)
 
 
 # Functions
-def keyboard_test():
+def keyboard_test() -> None:
   """Test keyboard using OS specific functions."""
   if PLATFORM == 'Linux':
     run_xev()
@@ -20,7 +20,7 @@ def keyboard_test():
     print_warning(f'Not supported under this OS: {PLATFORM}')
 
 
-def run_xev():
+def run_xev() -> None:
   """Test keyboard using xev."""
   LOG.info('Keyboard Test (xev)')
   cmd = ['xev', '-event', 'keyboard']

scripts/wk/hw/network.py

@@ -21,7 +21,7 @@ LOG = logging.getLogger(__name__)
 
 
 # Functions
-def network_test():
+def network_test() -> None:
   """Run network tests."""
   LOG.info('Network Test')
   try_and_print = TryAndPrint()

scripts/wk/hw/screensavers.py

@@ -14,7 +14,7 @@ LOG = logging.getLogger(__name__)
 
 
 # Functions
-def screensaver(name):
+def screensaver(name) -> None:
   """Show screensaver"""
   LOG.info('Screensaver (%s)', name)
   if name == 'matrix':

scripts/wk/hw/sensors.py

@@ -7,6 +7,7 @@ import pathlib
 import re
 
 from subprocess import CalledProcessError
+from typing import Any
 
 from wk.cfg.hw import CPU_CRITICAL_TEMP, SMC_IDS, TEMP_COLORS
 from wk.exe import run_program, start_thread
@@ -39,14 +40,14 @@ class Sensors():
     self.data = get_sensor_data()
     self.out_path = None
 
-  def clear_temps(self):
+  def clear_temps(self) -> None:
     """Clear saved temps but keep structure"""
     for adapters in self.data.values():
       for sources in adapters.values():
         for source_data in sources.values():
           source_data['Temps'] = []
 
-  def cpu_max_temp(self):
+  def cpu_max_temp(self) -> float:
     """Get max temp from any CPU source, returns float.
 
     NOTE: If no temps are found this returns zero.
@@ -64,7 +65,7 @@ class Sensors():
     # Done
     return max_temp
 
-  def cpu_reached_critical_temp(self):
+  def cpu_reached_critical_temp(self) -> bool:
     """Check if CPU reached CPU_CRITICAL_TEMP, returns bool."""
     for section, adapters in self.data.items():
       if not section.startswith('CPU'):
@@ -80,7 +81,8 @@ class Sensors():
     # Didn't return above so temps are within the threshold
     return False
 
-  def generate_report(self, *temp_labels, colored=True, only_cpu=False):
+  def generate_report(
+      self, *temp_labels, colored=True, only_cpu=False) -> list[str]:
     """Generate report based on given temp_labels, returns list."""
     report = []
 
@@ -117,7 +119,8 @@ class Sensors():
 
   def monitor_to_file(
       self, out_path, alt_max=None,
-      exit_on_thermal_limit=True, temp_labels=None, thermal_action=None):
+      exit_on_thermal_limit=True, temp_labels=None,
+      thermal_action=None) -> None:
     # pylint: disable=too-many-arguments
     """Write report to path every second until stopped.
 
@@ -151,7 +154,7 @@ class Sensors():
       # Sleep before next loop
       sleep(0.5)
 
-  def save_average_temps(self, temp_label, seconds=10):
+  def save_average_temps(self, temp_label, seconds=10) -> None:
     """Save average temps under temp_label over provided seconds.."""
     self.clear_temps()
 
@@ -177,7 +180,8 @@ class Sensors():
 
   def start_background_monitor(
       self, out_path, alt_max=None,
-      exit_on_thermal_limit=True, temp_labels=None, thermal_action=None):
+      exit_on_thermal_limit=True, temp_labels=None,
+      thermal_action=None) -> None:
     # pylint: disable=too-many-arguments
     """Start background thread to save report to file.
 
@@ -194,7 +198,7 @@ class Sensors():
         ),
       )
 
-  def stop_background_monitor(self):
+  def stop_background_monitor(self) -> None:
     """Stop background thread."""
     self.out_path.with_suffix('.stop').touch()
     self.background_thread.join()
@@ -203,14 +207,16 @@ class Sensors():
     self.background_thread = None
     self.out_path = None
 
-  def update_sensor_data(self, alt_max=None, exit_on_thermal_limit=True):
+  def update_sensor_data(
+      self, alt_max=None, exit_on_thermal_limit=True) -> None:
     """Update sensor data via OS-specific means."""
     if PLATFORM == 'Darwin':
       self.update_sensor_data_macos(alt_max, exit_on_thermal_limit)
     elif PLATFORM == 'Linux':
       self.update_sensor_data_linux(alt_max, exit_on_thermal_limit)
 
-  def update_sensor_data_linux(self, alt_max, exit_on_thermal_limit=True):
+  def update_sensor_data_linux(
+      self, alt_max, exit_on_thermal_limit=True) -> None:
     """Update sensor data via lm_sensors."""
     lm_sensor_data = get_sensor_data_lm()
     for section, adapters in self.data.items():
@@ -233,7 +239,8 @@ class Sensors():
             if source_data['Current'] >= CPU_CRITICAL_TEMP:
               raise ThermalLimitReachedError('CPU temps reached limit')
 
-  def update_sensor_data_macos(self, alt_max, exit_on_thermal_limit=True):
+  def update_sensor_data_macos(
+      self, alt_max, exit_on_thermal_limit=True) -> None:
     """Update sensor data via SMC."""
     for section, adapters in self.data.items():
       for sources in adapters.values():
@@ -261,7 +268,7 @@ class Sensors():
 
 
 # Functions
-def fix_sensor_name(name):
+def fix_sensor_name(name) -> str:
   """Cleanup sensor name, returns str."""
   name = re.sub(r'^(\w+)-(\w+)-(\w+)', r'\1 (\2 \3)', name, re.IGNORECASE)
   name = name.title()
@@ -280,7 +287,7 @@ def fix_sensor_name(name):
   return name
 
 
-def get_sensor_data():
+def get_sensor_data() -> dict[Any, Any]:
   """Get sensor data via OS-specific means, returns dict."""
   sensor_data = {}
   if PLATFORM == 'Darwin':
@@ -291,7 +298,7 @@ def get_sensor_data():
   return sensor_data
 
 
-def get_sensor_data_linux():
+def get_sensor_data_linux() -> dict[Any, Any]:
   """Get sensor data via lm_sensors, returns dict."""
   raw_lm_sensor_data = get_sensor_data_lm()
   sensor_data = {'CPUTemps': {}, 'Others': {}}
@@ -332,7 +339,7 @@ def get_sensor_data_linux():
   return sensor_data
 
 
-def get_sensor_data_lm():
+def get_sensor_data_lm() -> dict[Any, Any]:
   """Get raw sensor data via lm_sensors, returns dict."""
   raw_lm_sensor_data = {}
   cmd = ['sensors', '-j']
@@ -363,7 +370,7 @@ def get_sensor_data_lm():
   return raw_lm_sensor_data
 
 
-def get_sensor_data_macos():
+def get_sensor_data_macos() -> dict[Any, Any]:
   """Get sensor data via SMC, returns dict.
 
   NOTE: The data is structured like the lm_sensor data.
@@ -407,7 +414,7 @@ def get_sensor_data_macos():
   return sensor_data
 
 
-def get_temp_str(temp, colored=True):
+def get_temp_str(temp, colored=True) -> str:
   """Get colored string based on temp, returns str."""
   temp_color = None
 

scripts/wk/hw/smart.py

@@ -77,7 +77,7 @@ def check_attributes(dev, only_blocking=False) -> bool:
   return attributes_ok
 
 
-def check_self_test_results(test_obj, aborted=False):
+def check_self_test_results(test_obj, aborted=False) -> None:
   """Check SMART self-test results."""
   test_obj.report.append(color_string('Self-Test', 'BLUE'))
   if test_obj.disabled or test_obj.status == 'Denied':

scripts/wk/hw/surface_scan.py

@@ -20,7 +20,7 @@ LOG = logging.getLogger(__name__)
 
 
 # Functions
-def run_scan(test_obj, log_path):
+def run_scan(test_obj, log_path) -> None:
   """Run surface scan and handle exceptions."""
   block_size = '1024'
   dev = test_obj.dev

scripts/wk/hw/test.py

@@ -18,7 +18,7 @@ class Test:
   report: list[str] = field(init=False, default_factory=list)
   status: str = field(init=False, default='Pending')
 
-  def set_status(self, status):
+  def set_status(self, status) -> None:
     """Update status string."""
     if self.disabled:
       # Don't change status if disabled