From 20f91f01d1144fab22922a9a59a6fa3d0fe8f294 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Wed, 6 Apr 2022 19:08:23 -0600 Subject: [PATCH] Add type hints to hardware functions --- scripts/wk/hw/audio.py | 4 ++-- scripts/wk/hw/benchmark.py | 8 +++---- scripts/wk/hw/cpu.py | 19 ++++++++------- scripts/wk/hw/diags.py | 45 +++++++++++++++++------------------ scripts/wk/hw/keyboard.py | 4 ++-- scripts/wk/hw/network.py | 2 +- scripts/wk/hw/screensavers.py | 2 +- scripts/wk/hw/sensors.py | 41 ++++++++++++++++++------------- scripts/wk/hw/smart.py | 2 +- scripts/wk/hw/surface_scan.py | 2 +- scripts/wk/hw/test.py | 2 +- 11 files changed, 70 insertions(+), 61 deletions(-) diff --git a/scripts/wk/hw/audio.py b/scripts/wk/hw/audio.py index f6a984b5..484e2e24 100644 --- a/scripts/wk/hw/audio.py +++ b/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') diff --git a/scripts/wk/hw/benchmark.py b/scripts/wk/hw/benchmark.py index 78664192..2ee08577 100644 --- a/scripts/wk/hw/benchmark.py +++ b/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): - """Generate colored report using rate_list, returns list of str.""" +def check_io_results(test_obj, rate_list, graph_width) -> None: + """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': diff --git a/scripts/wk/hw/cpu.py b/scripts/wk/hw/cpu.py index 44329e7a..6fb1ec8f 100644 --- a/scripts/wk/hw/cpu.py +++ b/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: diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py index 051ab228..b4fccb13 100644 --- a/scripts/wk/hw/diags.py +++ b/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): - # pylint: disable=unused-argument + def fix_tmux_layout(self, forced=True) -> None: """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): - # pylint: disable=too-many-branches,too-many-statements +def disk_surface_scan(state, test_objects) -> None: + # 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) diff --git a/scripts/wk/hw/keyboard.py b/scripts/wk/hw/keyboard.py index 20a3db0a..68e2d0a6 100644 --- a/scripts/wk/hw/keyboard.py +++ b/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'] diff --git a/scripts/wk/hw/network.py b/scripts/wk/hw/network.py index 68ab30d9..700ebfea 100644 --- a/scripts/wk/hw/network.py +++ b/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() diff --git a/scripts/wk/hw/screensavers.py b/scripts/wk/hw/screensavers.py index 9bdfb719..4417777b 100644 --- a/scripts/wk/hw/screensavers.py +++ b/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': diff --git a/scripts/wk/hw/sensors.py b/scripts/wk/hw/sensors.py index 7785794d..6950f191 100644 --- a/scripts/wk/hw/sensors.py +++ b/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 diff --git a/scripts/wk/hw/smart.py b/scripts/wk/hw/smart.py index 6e1a1ab5..827d1c04 100644 --- a/scripts/wk/hw/smart.py +++ b/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': diff --git a/scripts/wk/hw/surface_scan.py b/scripts/wk/hw/surface_scan.py index 085bc52c..85d87f96 100644 --- a/scripts/wk/hw/surface_scan.py +++ b/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 diff --git a/scripts/wk/hw/test.py b/scripts/wk/hw/test.py index e9653d9a..ebed113a 100644 --- a/scripts/wk/hw/test.py +++ b/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