From 89fd64779263d6175ef69cb953ebdd17ab9c5b07 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sat, 1 Apr 2023 22:14:03 -0700 Subject: [PATCH 01/16] Split wk.std into debug, std, and ui sections --- scripts/auto_repairs.py | 6 +- scripts/auto_setup.py | 4 +- scripts/build-ufd | 2 +- scripts/build_kit_windows.py | 4 +- scripts/ddrescue-tui.py | 4 +- scripts/embedded_python_env.py | 2 +- scripts/hw-diags.py | 4 +- scripts/hw-sensors | 6 +- scripts/launch_sdio.py | 14 +- scripts/mount-all-volumes | 18 +- scripts/mount-backup-shares | 6 +- scripts/msword-search | 14 +- scripts/pyproject.toml | 1 + scripts/unmount-backup-shares | 6 +- scripts/upload-logs | 14 +- scripts/wk/__init__.py | 10 +- scripts/wk/clone/ddrescue.py | 223 ++++---- scripts/wk/debug.py | 121 ++++ scripts/wk/graph.py | 2 +- scripts/wk/hw/benchmark.py | 4 +- scripts/wk/hw/cpu.py | 6 +- scripts/wk/hw/diags.py | 100 ++-- scripts/wk/hw/disk.py | 3 +- scripts/wk/hw/keyboard.py | 3 +- scripts/wk/hw/network.py | 3 +- scripts/wk/hw/sensors.py | 3 +- scripts/wk/hw/smart.py | 3 +- scripts/wk/hw/surface_scan.py | 8 +- scripts/wk/hw/system.py | 8 +- scripts/wk/kit/build_win.py | 5 +- scripts/wk/kit/ufd.py | 90 +-- scripts/wk/log.py | 18 + scripts/wk/net.py | 3 +- scripts/wk/os/linux.py | 3 +- scripts/wk/os/win.py | 3 +- scripts/wk/repairs/win.py | 5 +- scripts/wk/setup/win.py | 5 +- scripts/wk/std.py | 976 --------------------------------- scripts/wk/ui/__init__.py | 3 + scripts/wk/ui/cli.py | 862 +++++++++++++++++++++++++++++ scripts/wk_debug.py | 8 +- 41 files changed, 1312 insertions(+), 1271 deletions(-) create mode 100644 scripts/wk/ui/__init__.py create mode 100644 scripts/wk/ui/cli.py diff --git a/scripts/auto_repairs.py b/scripts/auto_repairs.py index c01f1735..0615f115 100644 --- a/scripts/auto_repairs.py +++ b/scripts/auto_repairs.py @@ -5,7 +5,7 @@ import wk # Classes -REBOOT_STR = wk.std.color_string('Reboot', 'YELLOW') +REBOOT_STR = wk.ui.cli.color_string('Reboot', 'YELLOW') class MenuEntry(): """Simple class to allow cleaner code below.""" def __init__(self, name, function=None, selected=True, **kwargs): @@ -87,8 +87,8 @@ if __name__ == '__main__': try: wk.repairs.win.run_auto_repairs(BASE_MENUS, PRESETS) except KeyboardInterrupt: - wk.std.abort() + wk.ui.cli.abort() except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/auto_setup.py b/scripts/auto_setup.py index 071ab8a3..f76ee648 100644 --- a/scripts/auto_setup.py +++ b/scripts/auto_setup.py @@ -154,8 +154,8 @@ if __name__ == '__main__': try: wk.setup.win.run_auto_setup(BASE_MENUS, PRESETS) except KeyboardInterrupt: - wk.std.abort() + wk.ui.cli.abort() except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/build-ufd b/scripts/build-ufd index 8bb04f63..3dacb7f3 100755 --- a/scripts/build-ufd +++ b/scripts/build-ufd @@ -11,4 +11,4 @@ if __name__ == '__main__': except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/build_kit_windows.py b/scripts/build_kit_windows.py index e739014b..a9323fc1 100644 --- a/scripts/build_kit_windows.py +++ b/scripts/build_kit_windows.py @@ -8,8 +8,8 @@ if __name__ == '__main__': try: wk.kit.build.build_kit() except KeyboardInterrupt: - wk.std.abort() + wk.ui.cli.abort() except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/ddrescue-tui.py b/scripts/ddrescue-tui.py index 183595a3..8e273dda 100755 --- a/scripts/ddrescue-tui.py +++ b/scripts/ddrescue-tui.py @@ -12,7 +12,7 @@ if __name__ == '__main__': docopt(wk.clone.ddrescue.DOCSTRING) except SystemExit: print('') - wk.std.pause('Press Enter to exit...') + wk.ui.cli.pause('Press Enter to exit...') raise try: @@ -20,4 +20,4 @@ if __name__ == '__main__': except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/embedded_python_env.py b/scripts/embedded_python_env.py index 7ebf7b98..1c63099b 100644 --- a/scripts/embedded_python_env.py +++ b/scripts/embedded_python_env.py @@ -8,7 +8,7 @@ python.exe -i embedded_python_env.py import wk -wk.std.print_colored( +wk.ui.cli.print_colored( (wk.cfg.main.KIT_NAME_FULL, ': ', 'Debug Console'), ('GREEN', None, 'YELLOW'), sep='', diff --git a/scripts/hw-diags.py b/scripts/hw-diags.py index 13c456a3..e2d00a41 100755 --- a/scripts/hw-diags.py +++ b/scripts/hw-diags.py @@ -12,7 +12,7 @@ if __name__ == '__main__': docopt(wk.hw.diags.DOCSTRING) except SystemExit: print('') - wk.std.pause('Press Enter to exit...') + wk.ui.cli.pause('Press Enter to exit...') raise try: @@ -20,4 +20,4 @@ if __name__ == '__main__': except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/hw-sensors b/scripts/hw-sensors index 1f1393c9..55dd31d9 100755 --- a/scripts/hw-sensors +++ b/scripts/hw-sensors @@ -11,11 +11,11 @@ def main(): """Show sensor data on screen.""" sensors = wk.hw.sensors.Sensors() if platform.system() == 'Darwin': - wk.std.clear_screen() + wk.ui.cli.clear_screen() while True: print('\033[100A', end='') sensors.update_sensor_data() - wk.std.print_report(sensors.generate_report('Current', 'Max')) + wk.ui.cli.print_report(sensors.generate_report('Current', 'Max')) wk.std.sleep(1) elif platform.system() == 'Linux': proc = wk.exe.run_program(cmd=['mktemp']) @@ -43,4 +43,4 @@ if __name__ == '__main__': except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/launch_sdio.py b/scripts/launch_sdio.py index 03d4070c..8827c4d8 100644 --- a/scripts/launch_sdio.py +++ b/scripts/launch_sdio.py @@ -22,10 +22,10 @@ SDIO_REMOTE_PATH = wk.io.get_path_obj( # Functions def try_again(): """Ask to try again or quit.""" - if wk.std.ask(' Try again?'): + if wk.ui.cli.ask(' Try again?'): return True - if not wk.std.ask(' Use local version?'): - wk.std.abort() + if not wk.ui.cli.ask(' Use local version?'): + wk.ui.cli.abort() return False @@ -47,7 +47,7 @@ def use_network_sdio(): except KeyboardInterrupt: break except MOUNT_EXCEPTIONS as err: - wk.std.print_error(f' {err}') + wk.ui.cli.print_error(f' {err}') if not try_again(): break else: @@ -57,7 +57,7 @@ def use_network_sdio(): break # Failed to mount - wk.std.print_error(' Failed to mount server') + wk.ui.cli.print_error(' Failed to mount server') if not try_again(): break @@ -66,7 +66,7 @@ def use_network_sdio(): if __name__ == '__main__': - wk.std.set_title( + wk.ui.cli.set_title( f'{wk.cfg.main.KIT_NAME_FULL}: Snappy Driver Installer Origin Launcher', ) log_dir = wk.log.format_log_path(tool=True).parent @@ -76,7 +76,7 @@ if __name__ == '__main__': try: USE_NETWORK = use_network_sdio() except KeyboardInterrupt: - wk.std.abort() + wk.ui.cli.abort() # Run SDIO EXE_PATH = SDIO_LOCAL_PATH diff --git a/scripts/mount-all-volumes b/scripts/mount-all-volumes index 3adde487..53eb3a26 100755 --- a/scripts/mount-all-volumes +++ b/scripts/mount-all-volumes @@ -10,32 +10,32 @@ import wk # Functions def main(): """Mount all volumes and show results.""" - wk.std.print_standard(f'{wk.cfg.main.KIT_NAME_FULL}: Volume mount tool') - wk.std.print_standard(' ') + wk.ui.cli.print_standard(f'{wk.cfg.main.KIT_NAME_FULL}: Volume mount tool') + wk.ui.cli.print_standard(' ') # Mount volumes and get report - wk.std.print_standard('Mounting volumes...') + wk.ui.cli.print_standard('Mounting volumes...') wk.os.linux.mount_volumes() report = wk.os.linux.build_volume_report() # Show results - wk.std.print_info('Results') - wk.std.print_report(report) + wk.ui.cli.print_info('Results') + wk.ui.cli.print_report(report) # GUI mode if 'gui' in sys.argv: - wk.std.pause('Press Enter to exit...') + wk.ui.cli.pause('Press Enter to exit...') wk.exe.popen_program(['nohup', 'thunar', '/media']) if __name__ == '__main__': if wk.std.PLATFORM != 'Linux': os_name = wk.std.PLATFORM.replace('Darwin', 'macOS') - wk.std.print_error(f'This script is not supported under {os_name}.') - wk.std.abort() + wk.ui.cli.print_error(f'This script is not supported under {os_name}.') + wk.ui.cli.abort() try: main() except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/mount-backup-shares b/scripts/mount-backup-shares index d19ffaf6..bf528771 100755 --- a/scripts/mount-backup-shares +++ b/scripts/mount-backup-shares @@ -8,7 +8,7 @@ import wk # Functions def main(): """Attempt to mount backup shares and print report.""" - wk.std.print_info('Mounting Backup Shares') + wk.ui.cli.print_info('Mounting Backup Shares') report = wk.net.mount_backup_shares() for line in report: color = 'GREEN' @@ -17,7 +17,7 @@ def main(): color = 'RED' elif 'Already' in line: color = 'YELLOW' - print(wk.std.color_string(line, color)) + print(wk.ui.cli.color_string(line, color)) if __name__ == '__main__': @@ -26,4 +26,4 @@ if __name__ == '__main__': except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/msword-search b/scripts/msword-search index a5aac0b6..14062725 100755 --- a/scripts/msword-search +++ b/scripts/msword-search @@ -46,13 +46,13 @@ def scan_file(file_path, search): if __name__ == '__main__': try: # Prep - wk.std.clear_screen() + wk.ui.cli.clear_screen() terms = [re.sub(r'\s+', r'\s*', t) for t in sys.argv[1:]] search = '({})'.format('|'.join(terms)) if len(sys.argv) == 1: # Print usage - wk.std.print_standard(USAGE) + wk.ui.cli.print_standard(USAGE) else: matches = [] for entry in scan_for_docs(SCANDIR): @@ -60,20 +60,20 @@ if __name__ == '__main__': # Strip None values (i.e. non-matching entries) matches = [m for m in matches if m] if matches: - wk.std.print_success('Found {} {}:'.format( + wk.ui.cli.print_success('Found {} {}:'.format( len(matches), 'Matches' if len(matches) > 1 else 'Match')) for match in matches: - wk.std.print_standard(match) + wk.ui.cli.print_standard(match) else: - wk.std.print_error('No matches found.') + wk.ui.cli.print_error('No matches found.') # Done - wk.std.print_standard('\nDone.') + wk.ui.cli.print_standard('\nDone.') #pause("Press Enter to exit...") except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() # vim: sts=2 sw=2 ts=2 diff --git a/scripts/pyproject.toml b/scripts/pyproject.toml index 05f92df9..04adb66c 100644 --- a/scripts/pyproject.toml +++ b/scripts/pyproject.toml @@ -8,6 +8,7 @@ "wk/os/__init__.py" = ["F401"] "wk/repairs/__init__.py" = ["F401"] "wk/setup/__init__.py" = ["F401"] +"wk/ui/__init__.py" = ["F401"] # Long lines "wk/borrowed/acpi.py" = ["E501", "F841"] diff --git a/scripts/unmount-backup-shares b/scripts/unmount-backup-shares index e7375c08..57e7e880 100755 --- a/scripts/unmount-backup-shares +++ b/scripts/unmount-backup-shares @@ -8,14 +8,14 @@ import wk # Functions def main(): """Attempt to mount backup shares and print report.""" - wk.std.print_info('Unmounting Backup Shares') + wk.ui.cli.print_info('Unmounting Backup Shares') report = wk.net.unmount_backup_shares() for line in report: color = 'GREEN' line = f' {line}' if 'Not mounted' in line: color = 'YELLOW' - print(wk.std.color_string(line, color)) + print(wk.ui.cli.color_string(line, color)) if __name__ == '__main__': @@ -24,4 +24,4 @@ if __name__ == '__main__': except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/upload-logs b/scripts/upload-logs index 61ba67bf..6ea27f14 100755 --- a/scripts/upload-logs +++ b/scripts/upload-logs @@ -28,21 +28,21 @@ if PLATFORM not in ('macOS', 'Linux'): def main(): """Upload logs for review.""" lines = [] - try_and_print = wk.std.TryAndPrint() + try_and_print = wk.ui.cli.TryAndPrint() # Set log wk.log.update_log_path(dest_name='Upload-Logs', timestamp=True) # Instructions - wk.std.print_success(f'{wk.cfg.main.KIT_NAME_FULL}: Upload Logs') - wk.std.print_standard('') - wk.std.print_standard('Please state the reason for the review.') - wk.std.print_info(' End note with an empty line.') - wk.std.print_standard('') + wk.ui.cli.print_success(f'{wk.cfg.main.KIT_NAME_FULL}: Upload Logs') + wk.ui.cli.print_standard('') + wk.ui.cli.print_standard('Please state the reason for the review.') + wk.ui.cli.print_info(' End note with an empty line.') + wk.ui.cli.print_standard('') # Get reason note while True: - text = wk.std.input_text('> ') + text = wk.ui.cli.input_text('> ') if not text: lines.append('') break diff --git a/scripts/wk/__init__.py b/scripts/wk/__init__.py index e6877c23..5bf5d701 100644 --- a/scripts/wk/__init__.py +++ b/scripts/wk/__init__.py @@ -1,7 +1,7 @@ """WizardKit: wk module init""" # vim: sts=2 sw=2 ts=2 -from sys import version_info as version +from sys import stderr, version_info from . import cfg from . import clone @@ -18,20 +18,22 @@ from . import repairs from . import setup from . import std from . import tmux +from . import ui # Check env -if version < (3, 7): +if version_info < (3, 7): # Unsupported raise RuntimeError( - f'This package is unsupported on Python {version.major}.{version.minor}' + 'This package is unsupported on Python ' + f'{version_info.major}.{version_info.minor}' ) # Init try: log.start() except UserWarning as err: - std.print_warning(err) + print(err, file=stderr) if __name__ == '__main__': diff --git a/scripts/wk/clone/ddrescue.py b/scripts/wk/clone/ddrescue.py index 29667bd9..d0255360 100644 --- a/scripts/wk/clone/ddrescue.py +++ b/scripts/wk/clone/ddrescue.py @@ -33,6 +33,7 @@ from wk.hw.smart import ( smart_status_ok, update_smart_details, ) +from wk.ui import cli as ui # TODO: This is lazy # STATIC VARIABLES @@ -89,8 +90,8 @@ REGEX_REMAINING_TIME = re.compile( LOG = logging.getLogger(__name__) MENU_ACTIONS = ( 'Start', - f'Change settings {std.color_string("(experts only)", "YELLOW")}', - f'Detect drives {std.color_string("(experts only)", "YELLOW")}', + f'Change settings {ui.color_string("(experts only)", "YELLOW")}', + f'Detect drives {ui.color_string("(experts only)", "YELLOW")}', 'Quit') MENU_TOGGLES = { 'Auto continue (if recovery % over threshold)': True, @@ -296,7 +297,7 @@ class BlockPair(): # Check destination size if cloning if not self.destination.is_file() and dest_size < self.size: - std.print_error(f'Invalid destination: {self.destination}') + ui.print_error(f'Invalid destination: {self.destination}') raise std.GenericAbort() def set_initial_status(self): @@ -420,7 +421,7 @@ class State(): self.panes['Started'] = tmux.split_window( lines=cfg.ddrescue.TMUX_SIDE_WIDTH, target_id=self.panes['Source'], - text=std.color_string( + text=ui.color_string( ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], ['BLUE', None], sep='\n', @@ -442,7 +443,7 @@ class State(): settings = json.loads(_f.read()) except (OSError, json.JSONDecodeError) as err: LOG.error('Failed to load clone settings') - std.print_error('Invalid clone settings detected.') + ui.print_error('Invalid clone settings detected.') raise std.GenericAbort() from err # Check settings @@ -454,10 +455,10 @@ class State(): bail = False for key in ('model', 'serial'): if settings['Source'][key] != getattr(self.source, key): - std.print_error(f"Clone settings don't match source {key}") + ui.print_error(f"Clone settings don't match source {key}") bail = True if settings['Destination'][key] != getattr(self.destination, key): - std.print_error(f"Clone settings don't match destination {key}") + ui.print_error(f"Clone settings don't match destination {key}") bail = True if bail: raise std.GenericAbort() @@ -488,7 +489,7 @@ class State(): with open(settings_file, 'w', encoding='utf-8') as _f: json.dump(settings, _f) except OSError as err: - std.print_error('Failed to save clone settings') + ui.print_error('Failed to save clone settings') raise std.GenericAbort() from err def add_clone_block_pairs(self): @@ -522,7 +523,7 @@ class State(): # New run, use new settings file settings['Needs Format'] = True offset = 0 - user_choice = std.choice( + user_choice = ui.choice( ['G', 'M', 'S'], 'Format clone using GPT, MBR, or match Source type?', ) @@ -533,7 +534,7 @@ class State(): else: # Match source type settings['Table Type'] = get_table_type(self.source.path) - if std.ask('Create an empty Windows boot partition on the clone?'): + if ui.ask('Create an empty Windows boot partition on the clone?'): settings['Create Boot Partition'] = True offset = 2 if settings['Table Type'] == 'GPT' else 1 @@ -566,14 +567,14 @@ class State(): report = [] # Source - report.append(std.color_string('Source', 'GREEN')) + report.append(ui.color_string('Source', 'GREEN')) report.extend(build_object_report(self.source)) report.append(' ') # Destination - report.append(std.color_string('Destination', 'GREEN')) + report.append(ui.color_string('Destination', 'GREEN')) if self.mode == 'Clone': - report[-1] += std.color_string(' (ALL DATA WILL BE DELETED)', 'RED') + report[-1] += ui.color_string(' (ALL DATA WILL BE DELETED)', 'RED') report.extend(build_object_report(self.destination)) report.append(' ') @@ -581,12 +582,12 @@ class State(): # NOTE: The check for block_pairs is to limit this section # to the second confirmation if self.mode == 'Clone' and self.block_pairs: - report.append(std.color_string('WARNING', 'YELLOW')) + report.append(ui.color_string('WARNING', 'YELLOW')) report.append( 'All data will be deleted from the destination listed above.', ) report.append( - std.color_string( + ui.color_string( ['This is irreversible and will lead to', 'DATA LOSS.'], ['YELLOW', 'RED'], ), @@ -605,18 +606,18 @@ class State(): # Map dir if self.working_dir: - report.append(std.color_string('Map Save Directory', 'GREEN')) + report.append(ui.color_string('Map Save Directory', 'GREEN')) report.append(f'{self.working_dir}/') report.append(' ') if not fstype_is_ok(self.working_dir, map_dir=True): report.append( - std.color_string( + ui.color_string( 'Map file(s) are being saved to a non-recommended filesystem.', 'YELLOW', ), ) report.append( - std.color_string( + ui.color_string( ['This is strongly discouraged and may lead to', 'DATA LOSS'], [None, 'RED'], ), @@ -625,11 +626,11 @@ class State(): # Source part(s) selected if source_parts: - report.append(std.color_string('Source Part(s) selected', 'GREEN')) + report.append(ui.color_string('Source Part(s) selected', 'GREEN')) if self.source.path.samefile(source_parts[0].path): report.append('Whole Disk') else: - report.append(std.color_string(f'{"NAME":<9} SIZE', 'BLUE')) + report.append(ui.color_string(f'{"NAME":<9} SIZE', 'BLUE')) for part in source_parts: report.append( f'{part.path.name:<9} ' @@ -638,9 +639,9 @@ class State(): report.append(' ') # Prompt user - std.clear_screen() - std.print_report(report) - if not std.ask(prompt): + ui.clear_screen() + ui.print_report(report) + if not ui.ask(prompt): raise std.GenericAbort() def generate_report(self): @@ -661,7 +662,7 @@ class State(): error_size = self.get_error_size() error_size_str = std.bytes_to_string(error_size, decimals=2) if error_size > 0: - error_size_str = std.color_string(error_size_str, 'YELLOW') + error_size_str = ui.color_string(error_size_str, 'YELLOW') percent = self.get_percent_recovered() percent = format_status_string(percent, width=0) report.append(f'Overall rescued: {percent}, error size: {error_size_str}') @@ -673,7 +674,7 @@ class State(): error_size = pair.get_error_size() error_size_str = std.bytes_to_string(error_size, decimals=2) if error_size > 0: - error_size_str = std.color_string(error_size_str, 'YELLOW') + error_size_str = ui.color_string(error_size_str, 'YELLOW') pair_size = std.bytes_to_string(pair.size, decimals=2) percent = pair.get_percent_recovered() percent = format_status_string(percent, width=0) @@ -704,7 +705,7 @@ class State(): def init_recovery(self, docopt_args): """Select source/dest and set env.""" - std.clear_screen() + ui.clear_screen() source_parts = [] # Set log @@ -795,8 +796,8 @@ class State(): try: exe.run_program(cmd) except subprocess.CalledProcessError: - std.print_error('Failed to unmount source and/or destination') - std.abort() + ui.print_error('Failed to unmount source and/or destination') + ui.abort() # Prep destination if self.mode == 'Clone': @@ -928,7 +929,7 @@ class State(): check=False, ) if proc.returncode != 0: - std.print_error('Error(s) encoundtered while formatting destination') + ui.print_error('Error(s) encoundtered while formatting destination') raise std.GenericAbort() # Update settings @@ -969,13 +970,13 @@ class State(): # Check for critical errors if not smart_status_ok(self.destination): - std.print_error( + ui.print_error( f'Critical error(s) detected for: {self.destination.path}', ) # Check for minor errors if not check_attributes(self.destination, only_blocking=False): - std.print_warning( + ui.print_warning( f'Attribute error(s) detected for: {self.destination.path}', ) @@ -1026,7 +1027,7 @@ class State(): destination_size *= 1.05 error_msg = 'Not enough free space on the destination' if required_size > destination_size: - std.print_error(error_msg) + ui.print_error(error_msg) raise std.GenericAbort() def save_debug_reports(self): @@ -1037,7 +1038,7 @@ class State(): debug_dir.mkdir() # State (self) - std.save_pickles({'state': self}, debug_dir) + debug.save_pickles({'state': self}, debug_dir) with open(f'{debug_dir}/state.report', 'a', encoding='utf-8') as _f: _f.write('[Debug report]\n') _f.write('\n'.join(debug.generate_object_report(self))) @@ -1064,10 +1065,10 @@ class State(): width = cfg.ddrescue.TMUX_SIDE_WIDTH # Status - report.append(std.color_string(f'{"Status":^{width}}', 'BLUE')) + report.append(ui.color_string(f'{"Status":^{width}}', 'BLUE')) if 'NEEDS ATTENTION' in overall_status: report.append( - std.color_string(f'{overall_status:^{width}}', 'YELLOW_BLINK'), + ui.color_string(f'{overall_status:^{width}}', 'YELLOW_BLINK'), ) else: report.append(f'{overall_status:^{width}}') @@ -1077,12 +1078,12 @@ class State(): if self.block_pairs: total_rescued = self.get_rescued_size() percent = self.get_percent_recovered() - report.append(std.color_string('Overall Progress', 'BLUE')) + report.append(ui.color_string('Overall Progress', 'BLUE')) report.append( f'Rescued: {format_status_string(percent, width=width-9)}', ) report.append( - std.color_string( + ui.color_string( [f'{std.bytes_to_string(total_rescued, decimals=2):>{width}}'], [get_percent_color(percent)], ), @@ -1091,7 +1092,7 @@ class State(): # Block pair progress for pair in self.block_pairs: - report.append(std.color_string(pair.source, 'BLUE')) + report.append(ui.color_string(pair.source, 'BLUE')) for name, status in pair.status.items(): name = name.title() report.append( @@ -1103,9 +1104,9 @@ class State(): if overall_status in ('Active', 'NEEDS ATTENTION'): etoc = get_etoc() report.append(separator) - report.append(std.color_string('Estimated Pass Finish', 'BLUE')) + report.append(ui.color_string('Estimated Pass Finish', 'BLUE')) if overall_status == 'NEEDS ATTENTION' or etoc == 'N/A': - report.append(std.color_string('N/A', 'YELLOW')) + report.append(ui.color_string('N/A', 'YELLOW')) else: report.append(etoc) @@ -1166,7 +1167,7 @@ class State(): source_str = _format_string(self.source, width) tmux.respawn_pane( self.panes['Source'], - text=std.color_string( + text=ui.color_string( ['Source', '' if source_exists else ' (Missing)', '\n', source_str], ['BLUE', 'RED', None, None], sep='', @@ -1181,7 +1182,7 @@ class State(): percent=50, vertical=False, target_id=self.panes['Source'], - text=std.color_string( + text=ui.color_string( ['Destination', '' if dest_exists else ' (Missing)', '\n', dest_str], ['BLUE', 'RED', None, None], sep='', @@ -1195,7 +1196,7 @@ def build_block_pair_report(block_pairs, settings): report = [] notes = [] if block_pairs: - report.append(std.color_string('Block Pairs', 'GREEN')) + report.append(ui.color_string('Block Pairs', 'GREEN')) else: # Bail early return report @@ -1214,7 +1215,7 @@ def build_block_pair_report(block_pairs, settings): if settings: if not settings['First Run']: notes.append( - std.color_string( + ui.color_string( ['NOTE:', 'Clone settings loaded from previous run.'], ['BLUE', None], ), @@ -1222,14 +1223,14 @@ def build_block_pair_report(block_pairs, settings): if settings['Needs Format'] and settings['Table Type']: msg = f'Destination will be formatted using {settings["Table Type"]}' notes.append( - std.color_string( + ui.color_string( ['NOTE:', msg], ['BLUE', None], ), ) if any(pair.get_rescued_size() > 0 for pair in block_pairs): notes.append( - std.color_string( + ui.color_string( ['NOTE:', 'Resume data loaded from map file(s).'], ['BLUE', None], ), @@ -1311,12 +1312,12 @@ def build_directory_report(path): for line in proc.stdout.splitlines(): line = line.replace('\n', '') if 'FSTYPE' in line: - line = std.color_string(f'{"PATH":<{width}}{line}', 'BLUE') + line = ui.color_string(f'{"PATH":<{width}}{line}', 'BLUE') else: line = f'{path:<{width}}{line}' report.append(line) else: - report.append(std.color_string('PATH', 'BLUE')) + report.append(ui.color_string('PATH', 'BLUE')) report.append(str(path)) # Done @@ -1352,7 +1353,7 @@ def build_disk_report(dev): # Partition details report.append( - std.color_string( + ui.color_string( ( f'{"NAME":<{widths["name"]}}' f'{" " if dev.children else ""}' @@ -1397,8 +1398,8 @@ def build_disk_report(dev): def build_main_menu(): - """Build main menu, returns wk.std.Menu.""" - menu = std.Menu(title=std.color_string('ddrescue TUI: Main Menu', 'GREEN')) + """Build main menu, returns wk.ui.cli.Menu.""" + menu = ui.Menu(title=ui.color_string('ddrescue TUI: Main Menu', 'GREEN')) menu.separator = ' ' # Add actions, options, etc @@ -1429,23 +1430,23 @@ def build_object_report(obj): def build_settings_menu(silent=True): - """Build settings menu, returns wk.std.Menu.""" + """Build settings menu, returns wk.ui.cli.Menu.""" title_text = [ - std.color_string('ddrescue TUI: Expert Settings', 'GREEN'), + ui.color_string('ddrescue TUI: Expert Settings', 'GREEN'), ' ', - std.color_string( + ui.color_string( ['These settings can cause', 'MAJOR DAMAGE', 'to drives'], ['YELLOW', 'RED', 'YELLOW'], ), 'Please read the manual before making changes', ] - menu = std.Menu(title='\n'.join(title_text)) + menu = ui.Menu(title='\n'.join(title_text)) menu.separator = ' ' preset = 'Default' if not silent: # Ask which preset to use print(f'Available ddrescue presets: {" / ".join(SETTING_PRESETS)}') - preset = std.choice(SETTING_PRESETS, 'Please select a preset:') + preset = ui.choice(SETTING_PRESETS, 'Please select a preset:') # Fix selection for _p in SETTING_PRESETS: @@ -1494,7 +1495,7 @@ def build_sfdisk_partition_line(table_type, dev_path, size, details): # Safety Check if not dest_type: - std.print_error(f'Failed to determine partition type for: {dev_path}') + ui.print_error(f'Failed to determine partition type for: {dev_path}') raise std.GenericAbort() # Add extra details @@ -1576,7 +1577,7 @@ def format_status_string(status, width): # Add color if necessary if color: - status_str = std.color_string(status_str, color) + status_str = ui.color_string(status_str, color) # Done return status_str @@ -1698,8 +1699,8 @@ def get_object(path): # Child/Parent check if obj.parent: - std.print_warning(f'"{obj.path}" is a child device') - if std.ask(f'Use parent device "{obj.parent}" instead?'): + ui.print_warning(f'"{obj.path}" is a child device') + if ui.ask(f'Use parent device "{obj.parent}" instead?'): obj = hw_disk.Disk(obj.parent) elif path.is_dir(): obj = path @@ -1710,7 +1711,7 @@ def get_object(path): # Abort if obj not set if not obj: - std.print_error(f'Invalid source/dest path: {path}') + ui.print_error(f'Invalid source/dest path: {path}') raise std.GenericAbort() # Done @@ -1775,7 +1776,7 @@ def get_table_type(disk_path): # Check type if table_type not in ('GPT', 'MBR'): - std.print_error(f'Unsupported partition table type: {table_type}') + ui.print_error(f'Unsupported partition table type: {table_type}') raise std.GenericAbort() # Done @@ -1789,7 +1790,7 @@ def get_working_dir(mode, destination, force_local=False): # Set ticket ID while ticket_id is None: - ticket_id = std.input_text( + ticket_id = ui.input_text( prompt='Please enter ticket ID:', allow_empty_response=False, ) @@ -1802,12 +1803,12 @@ def get_working_dir(mode, destination, force_local=False): try: path = pathlib.Path(destination).resolve() except TypeError as err: - std.print_error(f'Invalid destination: {destination}') + ui.print_error(f'Invalid destination: {destination}') raise std.GenericAbort() from err if path.exists() and fstype_is_ok(path, map_dir=False): working_dir = path elif mode == 'Clone' and not force_local: - std.print_info('Mounting backup shares...') + ui.print_info('Mounting backup shares...') net.mount_backup_shares(read_write=True) for server in cfg.net.BACKUP_SERVERS: path = pathlib.Path( @@ -1851,11 +1852,11 @@ def is_missing_source_or_destination(state): if hasattr(item, 'path'): if not item.path.exists(): missing = True - std.print_error(f'{name} disappeared') + ui.print_error(f'{name} disappeared') elif hasattr(item, 'exists'): if not item.exists(): missing = True - std.print_error(f'{name} disappeared') + ui.print_error(f'{name} disappeared') else: LOG.error('Unknown %s type: %s', name, item) @@ -1887,7 +1888,7 @@ def source_or_destination_changed(state): # Done if changed: - std.print_error('Source and/or Destination changed') + ui.print_error('Source and/or Destination changed') return changed @@ -1910,7 +1911,7 @@ def main(): state.init_recovery(args) except (FileNotFoundError, std.GenericAbort): is_missing_source_or_destination(state) - std.abort() + ui.abort() # Show menu while True: @@ -1928,18 +1929,18 @@ def main(): # Detect drives if 'Detect drives' in selection[0]: - std.clear_screen() - std.print_warning(DETECT_DRIVES_NOTICE) - if std.ask('Are you sure you proceed?'): - std.print_standard('Forcing controllers to rescan for devices...') + ui.clear_screen() + ui.print_warning(DETECT_DRIVES_NOTICE) + if ui.ask('Are you sure you proceed?'): + ui.print_standard('Forcing controllers to rescan for devices...') cmd = 'echo "- - -" | sudo tee /sys/class/scsi_host/host*/scan' exe.run_program(cmd, check=False, shell=True) if source_or_destination_changed(state): - std.abort() + ui.abort() # Start recovery if 'Start' in selection: - std.clear_screen() + ui.clear_screen() run_recovery(state, main_menu, settings_menu, dry_run=args['--dry-run']) # Quit @@ -1949,14 +1950,14 @@ def main(): break # Recovey < 100% - std.print_warning('Recovery is less than 100%') - if std.ask('Are you sure you want to quit?'): + ui.print_warning('Recovery is less than 100%') + if ui.ask('Are you sure you want to quit?'): break # Save results to log LOG.info('') for line in state.generate_report(): - LOG.info(' %s', std.strip_colors(line)) + LOG.info(' %s', ui.strip_colors(line)) def mount_raw_image(path): @@ -1970,7 +1971,7 @@ def mount_raw_image(path): # Check if not loopback_path: - std.print_error(f'Failed to mount image: {path}') + ui.print_error(f'Failed to mount image: {path}') # Register unmount atexit atexit.register(unmount_loopback_device, loopback_path) @@ -2037,7 +2038,7 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): cmd = build_ddrescue_cmd(block_pair, pass_name, settings) poweroff_source_after_idle = True state.update_progress_pane('Active') - std.clear_screen() + ui.clear_screen() warning_message = '' def _poweroff_source_drive(idle_minutes): @@ -2056,8 +2057,8 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): return if i % 600 == 0 and i > 0: if i == 600: - std.print_standard(' ', flush=True) - std.print_warning( + ui.print_standard(' ', flush=True) + ui.print_warning( f'Powering off source in {int((idle_minutes*60-i)/60)} minutes...', ) std.sleep(5) @@ -2067,10 +2068,10 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): cmd = ['sudo', 'hdparm', '-Y', source_dev] proc = exe.run_program(cmd, check=False) if proc.returncode: - std.print_error(f'Failed to poweroff source {source_dev}') + ui.print_error(f'Failed to poweroff source {source_dev}') else: - std.print_warning(f'Powered off source {source_dev}') - std.print_standard( + ui.print_warning(f'Powered off source {source_dev}') + ui.print_standard( 'Press Enter to return to main menu...', end='', flush=True, ) @@ -2080,7 +2081,7 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): now = datetime.datetime.now(tz=TIMEZONE).strftime('%Y-%m-%d %H:%M %Z') with open(f'{state.log_dir}/smart.out', 'w', encoding='utf-8') as _f: _f.write( - std.color_string( + ui.color_string( ['SMART Attributes', f'Updated: {now}\n'], ['BLUE', 'YELLOW'], sep='\t\t', @@ -2113,7 +2114,7 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): if warning_message: # Error detected on destination, stop recovery exe.stop_process(proc) - std.print_error(warning_message) + ui.print_error(warning_message) break if _i % 60 == 0: @@ -2161,17 +2162,17 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): if warning_message: print(' ') print(' ') - std.print_error('DDRESCUE PROCESS HALTED') + ui.print_error('DDRESCUE PROCESS HALTED') print(' ') - std.print_warning(warning_message) + ui.print_warning(warning_message) # Needs attention? if str(proc.poll()) != '0': state.update_progress_pane('NEEDS ATTENTION') - std.pause('Press Enter to return to main menu...') + ui.pause('Press Enter to return to main menu...') # Stop source poweroff countdown - std.print_standard('Stopping device poweroff countdown...', flush=True) + ui.print_standard('Stopping device poweroff countdown...', flush=True) poweroff_source_after_idle = False poweroff_thread.join() @@ -2187,12 +2188,12 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True): # Bail early if is_missing_source_or_destination(state): - std.print_standard('') - std.pause('Press Enter to return to main menu...') + ui.print_standard('') + ui.pause('Press Enter to return to main menu...') return if source_or_destination_changed(state): - std.print_standard('') - std.abort() + ui.print_standard('') + ui.abort() # Get settings for name, details in main_menu.toggles.items(): @@ -2248,9 +2249,9 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True): # Show warning if nothing was done if not attempted_recovery: - std.print_warning('No actions performed') - std.print_standard(' ') - std.pause('Press Enter to return to main menu...') + ui.print_warning('No actions performed') + ui.print_standard(' ') + ui.pause('Press Enter to return to main menu...') # Done state.save_debug_reports() @@ -2260,10 +2261,10 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True): def select_disk(prompt, skip_disk=None): """Select disk from list, returns Disk().""" - std.print_info('Scanning disks...') + ui.print_info('Scanning disks...') disks = hw_disk.get_disks() - menu = std.Menu( - title=std.color_string(f'ddrescue TUI: {prompt} Selection', 'GREEN'), + menu = ui.Menu( + title=ui.color_string(f'ddrescue TUI: {prompt} Selection', 'GREEN'), ) menu.disabled_str = 'Already selected' menu.separator = ' ' @@ -2305,9 +2306,9 @@ def select_disk(prompt, skip_disk=None): def select_disk_parts(prompt, disk): """Select disk parts from list, returns list of Disk().""" - title = std.color_string('ddrescue TUI: Partition Selection', 'GREEN') + title = ui.color_string('ddrescue TUI: Partition Selection', 'GREEN') title += f'\n\nDisk: {disk.path} {disk.description}' - menu = std.Menu(title) + menu = ui.Menu(title) menu.separator = ' ' menu.add_action('All') menu.add_action('None') @@ -2356,7 +2357,7 @@ def select_disk_parts(prompt, disk): if not menu.options: menu.add_option(whole_disk_str, {'Selected': True, 'Path': disk.path}) menu.title += '\n\n' - menu.title += std.color_string(' No partitions detected.', 'YELLOW') + menu.title += ui.color_string(' No partitions detected.', 'YELLOW') # Get selection _select_parts(menu) @@ -2370,13 +2371,13 @@ def select_disk_parts(prompt, disk): if len(object_list) == len(disk.children): # NOTE: This is not true if the disk has no partitions msg = f'Preserve partition table and unused space in {prompt.lower()}?' - if std.ask(msg): + if ui.ask(msg): # Replace part list with whole disk obj object_list = [disk.path] # Convert object_list to hw_disk.Disk() objects print(' ') - std.print_info('Getting disk/partition details...') + ui.print_info('Getting disk/partition details...') object_list = [hw_disk.Disk(path) for path in object_list] # Done @@ -2386,8 +2387,8 @@ def select_disk_parts(prompt, disk): def select_path(prompt): """Select path, returns pathlib.Path.""" invalid = False - menu = std.Menu( - title=std.color_string(f'ddrescue TUI: {prompt} Path Selection', 'GREEN'), + menu = ui.Menu( + title=ui.color_string(f'ddrescue TUI: {prompt} Path Selection', 'GREEN'), ) menu.separator = ' ' menu.add_action('Quit') @@ -2400,7 +2401,7 @@ def select_path(prompt): if 'Current directory' in selection: path = os.getcwd() elif 'Enter manually' in selection: - path = std.input_text('Please enter path: ') + path = ui.input_text('Please enter path: ') elif 'Quit' in selection: raise std.GenericAbort() @@ -2410,7 +2411,7 @@ def select_path(prompt): except TypeError: invalid = True if invalid or not path.is_dir(): - std.print_error(f'Invalid path: {path}') + ui.print_error(f'Invalid path: {path}') raise std.GenericAbort() # Done @@ -2429,7 +2430,7 @@ def set_mode(docopt_args): # Ask user if necessary if not mode: - answer = std.choice(['C', 'I'], 'Are we cloning or imaging?') + answer = ui.choice(['C', 'I'], 'Are we cloning or imaging?') if answer == 'C': mode = 'Clone' else: diff --git a/scripts/wk/debug.py b/scripts/wk/debug.py index 11da914d..69f400d8 100644 --- a/scripts/wk/debug.py +++ b/scripts/wk/debug.py @@ -1,6 +1,21 @@ """WizardKit: Debug Functions""" # vim: sts=2 sw=2 ts=2 +import inspect +import logging +import lzma +import os +import pickle +import platform +import re +import socket +import sys +import time + +import requests + +from wk.cfg.net import CRASH_SERVER +from wk.log import get_log_filepath, get_root_logger_path # Classes class Debug(): @@ -10,11 +25,55 @@ class Debug(): # STATIC VARIABLES +LOG = logging.getLogger(__name__) DEBUG_CLASS = Debug() METHOD_TYPE = type(DEBUG_CLASS.method) # Functions +def generate_debug_report(): + """Generate debug report, returns str.""" + platform_function_list = ( + 'architecture', + 'machine', + 'platform', + 'python_version', + ) + report = [] + + # Logging data + log_path = get_log_filepath() + if log_path: + report.append('------ Start Log -------') + report.append('') + with open(log_path, 'r', encoding='utf-8') as log_file: + report.extend(log_file.read().splitlines()) + report.append('') + report.append('------- End Log --------') + + # System + report.append('--- Start debug info ---') + report.append('') + report.append('[System]') + report.append(f' {"FQDN":<24} {socket.getfqdn()}') + for func in platform_function_list: + func_name = func.replace('_', ' ').capitalize() + func_result = getattr(platform, func)() + report.append(f' {func_name:<24} {func_result}') + report.append(f' {"Python sys.argv":<24} {sys.argv}') + report.append('') + + # Environment + report.append('[Environment Variables]') + for key, value in sorted(os.environ.items()): + report.append(f' {key:<24} {value}') + report.append('') + + # Done + report.append('---- End debug info ----') + return '\n'.join(report) + + def generate_object_report(obj, indent=0): """Generate debug report for obj, returns list.""" report = [] @@ -46,5 +105,67 @@ def generate_object_report(obj, indent=0): return report +def save_pickles(obj_dict, out_path=None): + """Save dict of objects using pickle.""" + LOG.info('Saving pickles') + + # Set path + if not out_path: + out_path = get_root_logger_path() + out_path = out_path.parent.joinpath('../debug').resolve() + + # Save pickles + try: + for name, obj in obj_dict.copy().items(): + if name.startswith('__') or inspect.ismodule(obj): + continue + with open(f'{out_path}/{name}.pickle', 'wb') as _f: + pickle.dump(obj, _f, protocol=pickle.HIGHEST_PROTOCOL) + except Exception: + LOG.error('Failed to save all the pickles', exc_info=True) + + +def upload_debug_report(report, compress=True, reason='DEBUG'): + """Upload debug report to CRASH_SERVER as specified in wk.cfg.main.""" + LOG.info('Uploading debug report to %s', CRASH_SERVER.get('Name', '?')) + headers = CRASH_SERVER.get('Headers', {'X-Requested-With': 'XMLHttpRequest'}) + if compress: + headers['Content-Type'] = 'application/octet-stream' + + # Check if the required server details are available + if not all(CRASH_SERVER.get(key, False) for key in ('Name', 'Url', 'User')): + msg = 'Server details missing, aborting upload.' + print(msg) + raise RuntimeError(msg) + + # Set filename (based on the logging config if possible) + filename = 'Unknown' + log_path = get_log_filepath() + if log_path: + # Strip everything but the prefix + filename = re.sub(r'^(.*)_(\d{4}-\d{2}-\d{2}.*)', r'\1', log_path.name) + filename = f'{filename}_{reason}_{time.strftime("%Y-%m-%d_%H%M%S%z")}.log' + LOG.debug('filename: %s', filename) + + # Compress report + if compress: + filename += '.xz' + xz_report = lzma.compress(report.encode('utf8')) + + # Upload report + url = f'{CRASH_SERVER["Url"]}/{filename}' + response = requests.put( + url, + auth=(CRASH_SERVER['User'], CRASH_SERVER.get('Pass', '')), + data=xz_report if compress else report, + headers=headers, + timeout=60, + ) + + # Check response + if not response.ok: + raise RuntimeError('Failed to upload report') + + if __name__ == '__main__': print("This file is not meant to be called directly.") diff --git a/scripts/wk/graph.py b/scripts/wk/graph.py index 885cfe8e..fe717545 100644 --- a/scripts/wk/graph.py +++ b/scripts/wk/graph.py @@ -3,7 +3,7 @@ import logging -from wk.std import color_string +from wk.ui.cli import color_string # TODO: This is lazy # STATIC VARIABLES diff --git a/scripts/wk/hw/benchmark.py b/scripts/wk/hw/benchmark.py index 8aebf51d..b336da5c 100644 --- a/scripts/wk/hw/benchmark.py +++ b/scripts/wk/hw/benchmark.py @@ -21,8 +21,8 @@ from wk.cfg.hw import ( THRESH_SSD_MIN, ) from wk.exe import run_program -from wk.std import ( - PLATFORM, +from wk.std import PLATFORM +from wk.ui.cli import ( # TODO: This is lazy strip_colors, color_string, ) diff --git a/scripts/wk/hw/cpu.py b/scripts/wk/hw/cpu.py index 8458b918..826e0fb9 100644 --- a/scripts/wk/hw/cpu.py +++ b/scripts/wk/hw/cpu.py @@ -10,13 +10,13 @@ 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 -from wk.std import ( - PLATFORM, +from wk.std import PLATFORM +from wk.tmux import respawn_pane as tmux_respawn_pane +from wk.ui.cli import ( # TODO: This is lazy color_string, print_error, print_warning, ) -from wk.tmux import respawn_pane as tmux_respawn_pane # STATIC VARIABLES diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py index 7b7cb755..6c7ead59 100644 --- a/scripts/wk/hw/diags.py +++ b/scripts/wk/hw/diags.py @@ -25,6 +25,8 @@ from wk.hw.network import network_test from wk.hw.screensavers import screensaver from wk.hw.test import Test, TestGroup +from wk.ui import cli as ui # TODO: This is lazy + # STATIC VARIABLES DOCSTRING = f'''{cfg.main.KIT_NAME_FULL}: Hardware Diagnostics @@ -86,9 +88,9 @@ class State(): self.panes = {} self.system = None self.test_groups = [] - self.top_text = std.color_string('Hardware Diagnostics', 'GREEN') + self.top_text = ui.color_string('Hardware Diagnostics', 'GREEN') if test_mode: - self.top_text += std.color_string(' (Test Mode)', 'YELLOW') + self.top_text += ui.color_string(' (Test Mode)', 'YELLOW') # Init tmux and start a background process to maintain layout self.init_tmux() @@ -160,8 +162,8 @@ class State(): keep_history=False, timestamp=False, ) - std.clear_screen() - std.print_info('Initializing...') + ui.clear_screen() + ui.print_info('Initializing...') # Progress Pane self.update_progress_pane() @@ -226,7 +228,7 @@ class State(): self.panes['Started'] = tmux.split_window( lines=cfg.hw.TMUX_SIDE_WIDTH, target_id=self.panes['Top'], - text=std.color_string( + text=ui.color_string( ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], ['BLUE', None], sep='\n', @@ -247,7 +249,7 @@ class State(): debug_dir.mkdir() # State (self) - std.save_pickles({'state': self}, debug_dir) + debug.save_pickles({'state': self}, debug_dir) with open(f'{debug_dir}/state.report', 'a', encoding='utf-8') as _f: _f.write('\n'.join(debug.generate_object_report(self))) @@ -289,7 +291,7 @@ class State(): """Update 'Started' pane following clock sync.""" tmux.respawn_pane( pane_id=self.panes['Started'], - text=std.color_string( + text=ui.color_string( ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], ['BLUE', None], sep='\n', @@ -302,9 +304,9 @@ class State(): width = cfg.hw.TMUX_SIDE_WIDTH for group in self.test_groups: - report.append(std.color_string(group.name, 'BLUE')) + report.append(ui.color_string(group.name, 'BLUE')) for test in group.test_objects: - report.append(std.color_string( + report.append(ui.color_string( [test.label, f'{test.status:>{width-len(test.label)}}'], [None, STATUS_COLORS.get(test.status, None)], sep='', @@ -324,9 +326,9 @@ class State(): # Functions -def build_menu(cli_mode=False, quick_mode=False) -> std.Menu: - """Build main menu, returns wk.std.Menu.""" - menu = std.Menu(title=None) +def build_menu(cli_mode=False, quick_mode=False) -> ui.Menu: + """Build main menu, returns wk.ui.cli.Menu.""" + menu = ui.Menu(title=None) # Add actions, options, etc for action in MENU_ACTIONS: @@ -418,11 +420,11 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: state.layout['Current'] = {'height': 3, 'Check': True} # Get idle temps - std.print_standard('Saving idle temps...') + ui.print_standard('Saving idle temps...') sensors.save_average_temps(temp_label='Idle', seconds=5) # Stress CPU - std.print_info('Running stress test') + ui.print_info('Running stress test') hw_cpu.set_apple_fan_speed('max') proc_mprime = hw_cpu.start_mprime(state.log_dir, prime_log) @@ -443,14 +445,14 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: state.update_progress_pane() # Get cooldown temp - std.clear_screen() - std.print_standard('Letting CPU cooldown...') + ui.clear_screen() + ui.print_standard('Letting CPU cooldown...') std.sleep(5) - std.print_standard('Saving cooldown temps...') + ui.print_standard('Saving cooldown temps...') sensors.save_average_temps(temp_label='Cooldown', seconds=5) # Check Prime95 results - test_mprime_obj.report.append(std.color_string('Prime95', 'BLUE')) + test_mprime_obj.report.append(ui.color_string('Prime95', 'BLUE')) hw_cpu.check_mprime_results( test_obj=test_mprime_obj, working_dir=state.log_dir, ) @@ -461,10 +463,10 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: ) if run_sysbench: LOG.info('CPU Test (Sysbench)') - std.print_standard('Letting CPU cooldown more...') + ui.print_standard('Letting CPU cooldown more...') std.sleep(30) - std.clear_screen() - std.print_info('Running alternate stress test') + ui.clear_screen() + ui.print_info('Running alternate stress test') print('') proc_sysbench, filehandle_sysbench = hw_cpu.start_sysbench( sensors, @@ -490,7 +492,7 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: state.update_progress_pane() # Check Cooling results - test_cooling_obj.report.append(std.color_string('Temps', 'BLUE')) + test_cooling_obj.report.append(ui.color_string('Temps', 'BLUE')) hw_cpu.check_cooling_results(test_cooling_obj, sensors, run_sysbench) # Cleanup @@ -545,8 +547,8 @@ def disk_io_benchmark( continue # Start benchmark - std.clear_screen() - std.print_report(test.dev.generate_report()) + ui.clear_screen() + ui.print_report(test.dev.generate_report()) test.set_status('Working') test_log = f'{state.log_dir}/{test.dev.path.name}_benchmark.out' tmux.respawn_pane( @@ -563,12 +565,12 @@ def disk_io_benchmark( # Something went wrong LOG.error('%s', err) test.set_status('ERROR') - test.report.append(std.color_string(' Unknown Error', 'RED')) + test.report.append(ui.color_string(' Unknown Error', 'RED')) # Mark test(s) aborted if necessary if aborted: test.set_status('Aborted') - test.report.append(std.color_string(' Aborted', 'YELLOW')) + test.report.append(ui.color_string(' Aborted', 'YELLOW')) break # Update progress after each test @@ -594,7 +596,7 @@ def disk_self_test(state, test_objects, test_mode=False) -> None: state.update_top_pane( f'Disk self-test{"s" if len(test_objects) > 1 else ""}', ) - std.print_info(f'Starting self-test{"s" if len(test_objects) > 1 else ""}') + ui.print_info(f'Starting self-test{"s" if len(test_objects) > 1 else ""}') show_failed_attributes(state) for test in reversed(test_objects): if test.disabled: @@ -691,7 +693,7 @@ def disk_surface_scan(state, test_objects, test_mode=False) -> None: state.update_top_pane( f'Disk Surface Scan{"s" if len(test_objects) > 1 else ""}', ) - std.print_info( + ui.print_info( f'Starting disk surface scan{"s" if len(test_objects) > 1 else ""}', ) show_failed_attributes(state) @@ -731,7 +733,7 @@ def disk_surface_scan(state, test_objects, test_mode=False) -> None: for test in test_objects: if not (test.disabled or test.passed or test.failed): test.set_status('Aborted') - test.report.append(std.color_string(' Aborted', 'YELLOW')) + test.report.append(ui.color_string(' Aborted', 'YELLOW')) # Cleanup state.update_progress_pane() @@ -785,9 +787,9 @@ def main() -> None: try: action() except KeyboardInterrupt: - std.print_warning('Aborted.') - std.print_standard('') - std.pause('Press Enter to return to main menu...') + ui.print_warning('Aborted.') + ui.print_standard('') + ui.pause('Press Enter to return to main menu...') if 'Clock Sync' in selection: state.update_clock() @@ -852,8 +854,8 @@ def run_diags(state, menu, quick_mode=False, test_mode=False) -> None: # Just return if no tests were selected if not state.test_groups: - std.print_warning('No tests selected?') - std.pause() + ui.print_warning('No tests selected?') + ui.pause() return # Run tests @@ -864,7 +866,7 @@ def run_diags(state, menu, quick_mode=False, test_mode=False) -> None: args = [group.test_objects] if group.name == 'Disk I/O Benchmark': args.append(menu.toggles['Skip USB Benchmarks']['Selected']) - std.clear_screen() + ui.clear_screen() try: function(state, *args, test_mode=test_mode) except (KeyboardInterrupt, std.GenericAbort): @@ -891,25 +893,25 @@ def run_diags(state, menu, quick_mode=False, test_mode=False) -> None: state.save_debug_reports() atexit.unregister(state.save_debug_reports) if quick_mode: - std.pause('Press Enter to exit...') + ui.pause('Press Enter to exit...') else: - std.pause('Press Enter to return to main menu...') + ui.pause('Press Enter to return to main menu...') def show_failed_attributes(state) -> None: """Show failed attributes for all disks.""" for dev in state.disks: - std.print_colored([dev.name, dev.description], ['CYAN', None]) - std.print_report( + ui.print_colored([dev.name, dev.description], ['CYAN', None]) + ui.print_report( hw_smart.generate_attribute_report(dev, only_failed=True), ) - std.print_standard('') + ui.print_standard('') def show_results(state) -> None: """Show test results by device.""" std.sleep(0.5) - std.clear_screen() + ui.clear_screen() state.update_top_pane('Results') # CPU Tests @@ -917,22 +919,22 @@ def show_results(state) -> None: group.name for group in state.test_groups if 'CPU' in group.name ] if cpu_tests_enabled: - std.print_success('CPU:') - std.print_report(state.system.generate_report()) - std.print_standard(' ') + ui.print_success('CPU:') + ui.print_report(state.system.generate_report()) + ui.print_standard(' ') # Disk Tests disk_tests_enabled = [ group.name for group in state.test_groups if 'Disk' in group.name ] if disk_tests_enabled: - std.print_success(f'Disk{"s" if len(state.disks) > 1 else ""}:') + ui.print_success(f'Disk{"s" if len(state.disks) > 1 else ""}:') for disk in state.disks: - std.print_report(disk.generate_report()) - std.print_standard(' ') + ui.print_report(disk.generate_report()) + ui.print_standard(' ') if not state.disks: - std.print_warning('No devices') - std.print_standard(' ') + ui.print_warning('No devices') + ui.print_standard(' ') def sync_clock() -> None: diff --git a/scripts/wk/hw/disk.py b/scripts/wk/hw/disk.py index 3c7c727c..13bd821a 100644 --- a/scripts/wk/hw/disk.py +++ b/scripts/wk/hw/disk.py @@ -19,7 +19,8 @@ from wk.hw.smart import ( generate_attribute_report, get_known_disk_attributes, ) -from wk.std import PLATFORM, color_string, strip_colors +from wk.std import PLATFORM +from wk.ui.cli import color_string, strip_colors # TODO: This is lazy # STATIC VARIABLES diff --git a/scripts/wk/hw/keyboard.py b/scripts/wk/hw/keyboard.py index 68e2d0a6..686fd52b 100644 --- a/scripts/wk/hw/keyboard.py +++ b/scripts/wk/hw/keyboard.py @@ -4,7 +4,8 @@ import logging from wk.exe import run_program -from wk.std import PLATFORM, print_warning +from wk.std import PLATFORM +from wk.ui.cli import print_warning # TODO: This is lazy # STATIC VARIABLES diff --git a/scripts/wk/hw/network.py b/scripts/wk/hw/network.py index 700ebfea..29a7f8ba 100644 --- a/scripts/wk/hw/network.py +++ b/scripts/wk/hw/network.py @@ -9,7 +9,8 @@ from wk.net import ( show_valid_addresses, speedtest, ) -from wk.std import ( +from wk.ui.cli import ( + # TODO: This is lazy TryAndPrint, pause, print_warning, diff --git a/scripts/wk/hw/sensors.py b/scripts/wk/hw/sensors.py index 49d693f0..3ce1ded1 100644 --- a/scripts/wk/hw/sensors.py +++ b/scripts/wk/hw/sensors.py @@ -12,7 +12,8 @@ from typing import Any from wk.cfg.hw import CPU_CRITICAL_TEMP, SMC_IDS, TEMP_COLORS from wk.exe import run_program, start_thread from wk.io import non_clobber_path -from wk.std import PLATFORM, color_string, sleep +from wk.std import PLATFORM, sleep +from wk.ui.cli import color_string # TODO: This is lazy # STATIC VARIABLES diff --git a/scripts/wk/hw/smart.py b/scripts/wk/hw/smart.py index d9f47f7a..e56db1ea 100644 --- a/scripts/wk/hw/smart.py +++ b/scripts/wk/hw/smart.py @@ -18,7 +18,8 @@ from wk.cfg.hw import ( SMART_SELF_TEST_START_TIMEOUT_IN_SECONDS, ) from wk.exe import get_json_from_command, run_program -from wk.std import bytes_to_string, color_string, sleep +from wk.std import bytes_to_string, sleep +from wk.ui.cli import color_string # TODO: This is lazy # STATIC VARIABLES diff --git a/scripts/wk/hw/surface_scan.py b/scripts/wk/hw/surface_scan.py index 49afd25c..b2dc99c6 100644 --- a/scripts/wk/hw/surface_scan.py +++ b/scripts/wk/hw/surface_scan.py @@ -14,12 +14,8 @@ from wk.cfg.hw import ( TEST_MODE_BADBLOCKS_LIMIT, ) from wk.exe import run_program -from wk.std import ( - PLATFORM, - bytes_to_string, - color_string, - strip_colors, - ) +from wk.std import PLATFORM, bytes_to_string +from wk.ui.cli import color_string, strip_colors # TODO: This is lazy # STATIC VARIABLES diff --git a/scripts/wk/hw/system.py b/scripts/wk/hw/system.py index 8da2ed15..6634b326 100644 --- a/scripts/wk/hw/system.py +++ b/scripts/wk/hw/system.py @@ -12,12 +12,8 @@ from wk.cfg.hw import KNOWN_RAM_VENDOR_IDS from wk.cfg.python import DATACLASS_DECORATOR_KWARGS from wk.exe import get_json_from_command, run_program from wk.hw.test import Test -from wk.std import ( - PLATFORM, - bytes_to_string, - color_string, - string_to_bytes, - ) +from wk.std import PLATFORM, bytes_to_string, string_to_bytes +from wk.ui.cli import color_string # TODO: This is lazy # STATIC VARIABLES diff --git a/scripts/wk/kit/build_win.py b/scripts/wk/kit/build_win.py index ff377790..bb363430 100644 --- a/scripts/wk/kit/build_win.py +++ b/scripts/wk/kit/build_win.py @@ -21,8 +21,9 @@ from wk.kit.tools import ( get_tool_path, ) from wk.log import update_log_path -from wk.std import ( - GenericError, +from wk.std import GenericError +from wk.ui.cli import ( + # TODO: This is lazy TryAndPrint, clear_screen, pause, diff --git a/scripts/wk/kit/ufd.py b/scripts/wk/kit/ufd.py index cc5e9969..f0cfa767 100644 --- a/scripts/wk/kit/ufd.py +++ b/scripts/wk/kit/ufd.py @@ -10,7 +10,7 @@ from subprocess import CalledProcessError from collections import OrderedDict from docopt import docopt -from wk import io, log, std +from wk import io, log from wk.cfg.main import KIT_NAME_FULL, KIT_NAME_SHORT from wk.cfg.ufd import ( BOOT_ENTRIES, @@ -23,6 +23,8 @@ from wk.cfg.ufd import ( from wk.exe import get_json_from_command, run_program from wk.os import linux +from wk.ui import cli as ui # TODO: This is lazy + # STATIC VARIABLES DOCSTRING = '''WizardKit: Build UFD @@ -93,10 +95,10 @@ def build_ufd(): if args['--debug']: log.enable_debug_mode() if args['--update'] and args['EXTRA_IMAGES']: - std.print_warning('Extra images are ignored when updating') + ui.print_warning('Extra images are ignored when updating') args['EXTRA_IMAGES'] = [] log.update_log_path(dest_name='build-ufd', timestamp=True) - try_print = std.TryAndPrint() + try_print = ui.TryAndPrint() try_print.add_error('FileNotFoundError') try_print.catch_all = False try_print.indent = 2 @@ -104,9 +106,9 @@ def build_ufd(): try_print.width = 64 # Show header - std.print_success(KIT_NAME_FULL) - std.print_warning('UFD Build Tool') - std.print_warning(' ') + ui.print_success(KIT_NAME_FULL) + ui.print_warning('UFD Build Tool') + ui.print_warning(' ') # Verify selections ufd_dev = verify_ufd(args['--ufd-device']) @@ -118,7 +120,7 @@ def build_ufd(): # Prep UFD if not args['--update']: - std.print_info('Prep UFD') + ui.print_info('Prep UFD') try_print.run( message='Zeroing first 64MiB...', function=zero_device, @@ -170,8 +172,8 @@ def build_ufd(): ) # Copy sources - std.print_standard(' ') - std.print_info('Copy Sources') + ui.print_standard(' ') + ui.print_info('Copy Sources') try_print.run( 'Copying Memtest86...', io.recursive_copy, '/usr/share/memtest86-efi/', '/mnt/UFD/EFI/Memtest86/', overwrite=True, @@ -187,8 +189,8 @@ def build_ufd(): # Apply extra images if not args['--update']: - std.print_standard(' ') - std.print_info('Apply Extra Images') + ui.print_standard(' ') + ui.print_info('Apply Extra Images') for part_num, image_path in enumerate(extra_images): try_print.run( message=f'Applying {image_path.name}...', @@ -203,8 +205,8 @@ def build_ufd(): _f.write('\n'.join([image.name for image in extra_images])) # Update boot entries - std.print_standard(' ') - std.print_info('Boot Setup') + ui.print_standard(' ') + ui.print_info('Boot Setup') try_print.run( message='Updating boot entries...', function=update_boot_entries, @@ -235,8 +237,8 @@ def build_ufd(): ) # Hide items - std.print_standard(' ') - std.print_info('Final Touches') + ui.print_standard(' ') + ui.print_info('Final Touches') try_print.run( message='Hiding items...', function=hide_items, @@ -245,30 +247,30 @@ def build_ufd(): ) # Done - std.print_standard('\nDone.') + ui.print_standard('\nDone.') if not args['--force']: - std.pause('Press Enter to exit...') + ui.pause('Press Enter to exit...') def confirm_selections(update=False): """Ask tech to confirm selections, twice if necessary.""" - if not std.ask('Is the above information correct?'): - std.abort() + if not ui.ask('Is the above information correct?'): + ui.abort() # Safety check if not update: - std.print_standard(' ') - std.print_warning('SAFETY CHECK') - std.print_standard( + ui.print_standard(' ') + ui.print_warning('SAFETY CHECK') + ui.print_standard( 'All data will be DELETED from the disk and partition(s) listed above.') - std.print_colored( + ui.print_colored( ['This is irreversible and will lead to', 'DATA LOSS'], [None, 'RED'], ) - if not std.ask('Asking again to confirm, is this correct?'): - std.abort() + if not ui.ask('Asking again to confirm, is this correct?'): + ui.abort() - std.print_standard(' ') + ui.print_standard(' ') def copy_source(source, items, overwrite=False): @@ -481,12 +483,12 @@ def show_selections(args, sources, ufd_dev, ufd_sources, extra_images): """Show selections including non-specified options.""" # Sources - std.print_info('Sources') + ui.print_info('Sources') for label in ufd_sources.keys(): if label in sources: - std.print_standard(f' {label+":":<18} {sources[label]}') + ui.print_standard(f' {label+":":<18} {sources[label]}') else: - std.print_colored( + ui.print_colored( [f' {label+":":<18}', 'Not Specified'], [None, 'YELLOW'], ) @@ -498,15 +500,15 @@ def show_selections(args, sources, ufd_dev, ufd_sources, extra_images): print(f' {" ":<18} {image}') # Destination - std.print_standard(' ') - std.print_info('Destination') + ui.print_standard(' ') + ui.print_info('Destination') cmd = [ 'lsblk', '--nodeps', '--noheadings', '--paths', '--output', 'NAME,FSTYPE,TRAN,SIZE,VENDOR,MODEL,SERIAL', ufd_dev, ] proc = run_program(cmd, check=False) - std.print_standard(proc.stdout.strip()) + ui.print_standard(proc.stdout.strip()) cmd = [ 'lsblk', '--noheadings', '--paths', '--output', 'NAME,SIZE,FSTYPE,LABEL,MOUNTPOINT', @@ -514,14 +516,14 @@ def show_selections(args, sources, ufd_dev, ufd_sources, extra_images): ] proc = run_program(cmd, check=False) for line in proc.stdout.splitlines()[1:]: - std.print_standard(line) + ui.print_standard(line) # Notes if args['--update']: - std.print_warning('Updating kit in-place') + ui.print_warning('Updating kit in-place') elif args['--use-mbr']: - std.print_warning('Formatting using legacy MBR') - std.print_standard(' ') + ui.print_warning('Formatting using legacy MBR') + ui.print_standard(' ') def update_boot_entries(ufd_dev, images=None): @@ -621,11 +623,11 @@ def verify_sources(args, ufd_sources): try: s_path_obj = io.case_insensitive_path(s_path) except FileNotFoundError: - std.print_error(f'ERROR: {label} not found: {s_path}') - std.abort() + ui.print_error(f'ERROR: {label} not found: {s_path}') + ui.abort() if not is_valid_path(s_path_obj, data['Type']): - std.print_error(f'ERROR: Invalid {label} source: {s_path}') - std.abort() + ui.print_error(f'ERROR: Invalid {label} source: {s_path}') + ui.abort() sources[label] = s_path_obj return sources @@ -638,12 +640,12 @@ def verify_ufd(dev_path): try: ufd_dev = io.case_insensitive_path(dev_path) except FileNotFoundError: - std.print_error(f'ERROR: UFD device not found: {dev_path}') - std.abort() + ui.print_error(f'ERROR: UFD device not found: {dev_path}') + ui.abort() if not is_valid_path(ufd_dev, 'UFD'): - std.print_error(f'ERROR: Invalid UFD device: {ufd_dev}') - std.abort() + ui.print_error(f'ERROR: Invalid UFD device: {ufd_dev}') + ui.abort() return ufd_dev diff --git a/scripts/wk/log.py b/scripts/wk/log.py index b69afb1f..0b3e7e06 100644 --- a/scripts/wk/log.py +++ b/scripts/wk/log.py @@ -61,6 +61,24 @@ def format_log_path( return log_path +def get_log_filepath(): + """Get the log filepath from the root logger, returns pathlib.Path obj. + + NOTE: This will use the first handler baseFilename it finds (if any). + """ + log_filepath = None + root_logger = logging.getLogger() + + # Check handlers + for handler in root_logger.handlers: + if hasattr(handler, 'baseFilename'): + log_filepath = pathlib.Path(handler.baseFilename).resolve() + break + + # Done + return log_filepath + + def get_root_logger_path(): """Get path to log file from root logger, returns pathlib.Path obj.""" log_path = None diff --git a/scripts/wk/net.py b/scripts/wk/net.py index 3e9382a9..1c4c3383 100644 --- a/scripts/wk/net.py +++ b/scripts/wk/net.py @@ -8,9 +8,10 @@ import re import psutil from wk.exe import get_json_from_command, run_program -from wk.std import PLATFORM, GenericError, show_data +from wk.std import PLATFORM, GenericError from wk.cfg.net import BACKUP_SERVERS +from wk.ui.cli import show_data # TODO: This is lazy # REGEX diff --git a/scripts/wk/os/linux.py b/scripts/wk/os/linux.py index 481e687c..1c134f78 100644 --- a/scripts/wk/os/linux.py +++ b/scripts/wk/os/linux.py @@ -10,7 +10,8 @@ import subprocess from wk.cfg.hw import VOLUME_FAILURE_THRESHOLD, VOLUME_WARNING_THRESHOLD from wk.exe import get_json_from_command, popen_program, run_program from wk.log import format_log_path -from wk.std import bytes_to_string, color_string +from wk.std import bytes_to_string +from wk.ui.cli import color_string # TODO: This is lazy # STATIC VARIABLES diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py index 3b5c80e3..bdf424ad 100644 --- a/scripts/wk/os/win.py +++ b/scripts/wk/os/win.py @@ -30,10 +30,9 @@ from wk.std import ( GenericError, GenericWarning, bytes_to_string, - color_string, - input_text, sleep, ) +from wk.ui.cli import color_string, input_text # TODO: This is lazy # STATIC VARIABLES diff --git a/scripts/wk/repairs/win.py b/scripts/wk/repairs/win.py index 938fb3b3..f468be33 100644 --- a/scripts/wk/repairs/win.py +++ b/scripts/wk/repairs/win.py @@ -58,6 +58,10 @@ from wk.os.win import ( from wk.std import ( GenericError, GenericWarning, + sleep, + ) +from wk.ui.cli import ( + # TODO: This is lazy Menu, TryAndPrint, abort, @@ -70,7 +74,6 @@ from wk.std import ( print_warning, set_title, show_data, - sleep, strip_colors, ) diff --git a/scripts/wk/setup/win.py b/scripts/wk/setup/win.py index 362c7b1a..ad388192 100644 --- a/scripts/wk/setup/win.py +++ b/scripts/wk/setup/win.py @@ -60,6 +60,10 @@ from wk.repairs.win import ( from wk.std import ( GenericError, GenericWarning, + sleep, + ) +from wk.ui.cli import ( + # TODO: This is lazy Menu, TryAndPrint, abort, @@ -73,7 +77,6 @@ from wk.std import ( print_warning, set_title, show_data, - sleep, strip_colors, ) diff --git a/scripts/wk/std.py b/scripts/wk/std.py index 65fd69f0..cea6b2e1 100644 --- a/scripts/wk/std.py +++ b/scripts/wk/std.py @@ -1,55 +1,13 @@ """WizardKit: Standard Functions""" # vim: sts=2 sw=2 ts=2 -import inspect -import itertools import logging -import lzma -import os -import pathlib -import pickle import platform import re -import socket -import subprocess -import sys import time -import traceback - -from collections import OrderedDict - -try: - from functools import cache -except ImportError: - # Assuming Python is < 3.9 - from functools import lru_cache as cache - -import requests - -from wk.cfg.main import ( - ENABLED_UPLOAD_DATA, - INDENT, - SUPPORT_MESSAGE, - WIDTH, - ) -from wk.cfg.net import CRASH_SERVER -from wk.log import get_root_logger_path # STATIC VARIABLES -COLORS = { - 'CLEAR': '\033[0m', - 'RED': '\033[31m', - 'RED_BLINK': '\033[31;5m', - 'ORANGE': '\033[31;1m', - 'ORANGE_RED': '\033[1;31;41m', - 'GREEN': '\033[32m', - 'YELLOW': '\033[33m', - 'YELLOW_BLINK': '\033[33;5m', - 'BLUE': '\033[34m', - 'PURPLE': '\033[35m', - 'CYAN': '\033[36m', - } LOG = logging.getLogger(__name__) PLATFORM = platform.system() REGEX_SIZE_STRING = re.compile( @@ -72,574 +30,7 @@ class GenericWarning(Exception): """ -# Classes -class Menu(): - """Object for tracking menu specific data and methods. - - Menu items are added to an OrderedDict so the order is preserved. - - ASSUMPTIONS: - 1. All entry names are unique. - 2. All action entry names start with different letters. - """ - def __init__(self, title='[Untitled Menu]'): - self.actions = OrderedDict() - self.options = OrderedDict() - self.sets = OrderedDict() - self.toggles = OrderedDict() - self.disabled_str = 'Disabled' - self.separator = '─' - self.title = title - - def _generate_menu_text(self): - """Generate menu text, returns str.""" - separator_string = self._get_separator_string() - menu_lines = [self.title, separator_string] if self.title else [] - - # Sets & toggles - for section in (self.sets, self.toggles): - for details in section.values(): - if details.get('Hidden', False): - continue - if details.get('Separator', False): - menu_lines.append(separator_string) - menu_lines.append(details['Display Name']) - if self.sets or self.toggles: - menu_lines.append(separator_string) - - # Options - for details in self.options.values(): - if details.get('Hidden', False): - continue - if details.get('Separator', False): - menu_lines.append(separator_string) - menu_lines.append(details['Display Name']) - if self.options: - menu_lines.append(separator_string) - - # Actions - for details in self.actions.values(): - if details.get('Hidden', False): - continue - if details.get('Separator', False): - menu_lines.append(separator_string) - menu_lines.append(details['Display Name']) - - # Show menu - menu_lines.append('') - menu_lines = [str(line) for line in menu_lines] - return '\n'.join(menu_lines) - - def _get_display_name( - self, name, details, - index=None, no_checkboxes=True, setting_item=False): - """Format display name based on details and args, returns str.""" - disabled = details.get('Disabled', False) - if setting_item and not details['Selected']: - # Display item in YELLOW - disabled = True - checkmark = '*' - if 'DISPLAY' in os.environ or PLATFORM == 'Darwin': - checkmark = '✓' - display_name = f'{index if index else name[:1].upper()}: ' - if not (index and index >= 10): - display_name = f' {display_name}' - if setting_item and 'Value' in details: - name = f'{name} = {details["Value"]}' - - # Add enabled status if necessary - if not no_checkboxes: - display_name += f'[{checkmark if details["Selected"] else " "}] ' - - # Add name - if disabled: - display_name += color_string(f'{name} ({self.disabled_str})', 'YELLOW') - else: - display_name += name - - # Done - return display_name - - def _get_separator_string(self): - """Format separator length based on name lengths, returns str.""" - separator_length = 0 - - # Check title line(s) - if self.title: - for line in self.title.split('\n'): - separator_length = max(separator_length, len(strip_colors(line))) - - # Loop over all item names - for section in (self.actions, self.options, self.sets, self.toggles): - for details in section.values(): - if details.get('Hidden', False): - # Skip hidden lines - continue - line = strip_colors(details['Display Name']) - separator_length = max(separator_length, len(line)) - separator_length += 1 - - # Done - return self.separator * separator_length - - def _get_valid_answers(self): - """Get valid answers based on menu items, returns list.""" - valid_answers = [] - - # Numbered items - index = 0 - for section in (self.sets, self.toggles, self.options): - for details in section.values(): - if details.get('Hidden', False): - # Don't increment index or add to valid_answers - continue - index += 1 - if not details.get('Disabled', False): - valid_answers.append(str(index)) - - # Action items - for name, details in self.actions.items(): - if not details.get('Disabled', False): - valid_answers.append(name[:1].upper()) - - # Done - return valid_answers - - def _resolve_selection(self, selection): - """Get menu item based on user selection, returns tuple.""" - offset = 1 - resolved_selection = None - if selection.isnumeric(): - # Enumerate over numbered entries - entries = [ - *self.sets.items(), - *self.toggles.items(), - *self.options.items(), - ] - for _i, details in enumerate(entries): - if details[1].get('Hidden', False): - offset -= 1 - elif str(_i+offset) == selection: - resolved_selection = (details) - break - else: - # Just check actions - for action, details in self.actions.items(): - if action.lower().startswith(selection.lower()): - resolved_selection = (action, details) - break - - # Done - return resolved_selection - - def _update(self, single_selection=True, settings_mode=False): - """Update menu items in preparation for printing to screen.""" - index = 0 - - # Fix selection status for sets - for set_details in self.sets.values(): - set_selected = True - set_targets = set_details['Targets'] - for option, option_details in self.options.items(): - if option in set_targets and not option_details['Selected']: - set_selected = False - elif option not in set_targets and option_details['Selected']: - set_selected = False - set_details['Selected'] = set_selected - - # Numbered sections - for section in (self.sets, self.toggles, self.options): - for name, details in section.items(): - if details.get('Hidden', False): - # Skip hidden lines and don't increment index - continue - index += 1 - details['Display Name'] = self._get_display_name( - name, - details, - index=index, - no_checkboxes=single_selection, - setting_item=settings_mode, - ) - - # Actions - for name, details in self.actions.items(): - details['Display Name'] = self._get_display_name( - name, - details, - no_checkboxes=True, - ) - - def _update_entry_selection_status(self, entry, toggle=True, status=None): - """Update entry selection status either directly or by toggling.""" - if entry in self.sets: - # Update targets not the set itself - new_status = not self.sets[entry]['Selected'] if toggle else status - targets = self.sets[entry]['Targets'] - self._update_set_selection_status(targets, new_status) - for section in (self.toggles, self.options, self.actions): - if entry in section: - if toggle: - section[entry]['Selected'] = not section[entry]['Selected'] - else: - section[entry]['Selected'] = status - - def _update_set_selection_status(self, targets, status): - """Select or deselect options based on targets and status.""" - for option, details in self.options.items(): - # If (new) status is True and this option is a target then select - # Otherwise deselect - details['Selected'] = status and option in targets - - def _user_select(self, prompt): - """Show menu and select an entry, returns str.""" - menu_text = self._generate_menu_text() - valid_answers = self._get_valid_answers() - - # Menu loop - while True: - clear_screen() - print(menu_text) - sleep(0.01) - answer = input_text(prompt).strip() - if answer.upper() in valid_answers: - break - - # Done - return answer - - def add_action(self, name, details=None): - """Add action to menu.""" - details = details if details else {} - details['Selected'] = details.get('Selected', False) - self.actions[name] = details - - def add_option(self, name, details=None): - """Add option to menu.""" - details = details if details else {} - details['Selected'] = details.get('Selected', False) - self.options[name] = details - - def add_set(self, name, details=None): - """Add set to menu.""" - details = details if details else {} - details['Selected'] = details.get('Selected', False) - - # Safety check - if 'Targets' not in details: - raise KeyError('Menu set has no targets') - - # Add set - self.sets[name] = details - - def add_toggle(self, name, details=None): - """Add toggle to menu.""" - details = details if details else {} - details['Selected'] = details.get('Selected', False) - self.toggles[name] = details - - def advanced_select(self, prompt='Please make a selection: '): - """Display menu and make multiple selections, returns tuple. - - NOTE: Menu is displayed until an action entry is selected. - """ - while True: - self._update(single_selection=False) - user_selection = self._user_select(prompt) - selected_entry = self._resolve_selection(user_selection) - if user_selection.isnumeric(): - # Update selection(s) - self._update_entry_selection_status(selected_entry[0]) - else: - # Action selected - break - - # Done - return selected_entry - - def settings_select(self, prompt='Please make a selection: '): - """Display menu and make multiple selections, returns tuple. - - NOTE: Menu is displayed until an action entry is selected. - """ - choice_kwargs = { - 'choices': ['T', 'C'], - 'prompt': 'Toggle or change value?', - } - - while True: - self._update(single_selection=True, settings_mode=True) - user_selection = self._user_select(prompt) - selected_entry = self._resolve_selection(user_selection) - if user_selection.isnumeric(): - if 'Value' in selected_entry[-1] and choice(**choice_kwargs) == 'C': - # Change - selected_entry[-1]['Value'] = input_text('Enter new value: ') - else: - # Toggle - self._update_entry_selection_status(selected_entry[0]) - else: - # Action selected - break - - # Done - return selected_entry - - def simple_select(self, prompt='Please make a selection: ', update=True): - """Display menu and make a single selection, returns tuple.""" - if update: - self._update() - user_selection = self._user_select(prompt) - return self._resolve_selection(user_selection) - - def update(self): - """Update menu with default settings.""" - self._update() - - -class TryAndPrint(): - """Object used to standardize running functions and returning the result. - - The errors and warning attributes are used to allow fine-tuned results - based on exception names. - """ - def __init__(self, msg_bad='FAILED', msg_good='SUCCESS'): - self.catch_all = True - self.indent = INDENT - self.list_errors = ['GenericError'] - self.list_warnings = ['GenericWarning'] - self.msg_bad = msg_bad - self.msg_good = msg_good - self.verbose = False - self.width = WIDTH - - def _format_exception_message(self, _exception): - """Format using the exception's args or name, returns str.""" - LOG.debug( - 'Formatting exception: %s, %s', - _exception.__class__.__name__, - _exception, - ) - message = '' - - # Format message string from _exception - try: - if isinstance(_exception, subprocess.CalledProcessError): - message = _exception.stderr - if not isinstance(message, str): - message = message.decode('utf-8') - message = message.strip() - elif isinstance(_exception, ZeroDivisionError): - # Skip and just use exception name below - pass - else: - message = str(_exception) - except Exception: - # Just use the exception name instead - pass - - # Prepend exception name - if _exception.__class__.__name__ not in ('GenericError', 'GenericWarning'): - try: - message = f'{_exception.__class__.__name__}: {message}' - except Exception: - message = f'UNKNOWN ERROR: {message}' - - # Fix multi-line messages - if '\n' in message: - try: - lines = [ - f'{" "*(self.indent+self.width)}{line.strip()}' - for line in message.splitlines() if line.strip() - ] - lines[0] = lines[0].strip() - message = '\n'.join(lines) - except Exception: - pass - - # Done - return message - - def _format_function_output(self, output, msg_good): - """Format function output for use in try_and_print(), returns str.""" - LOG.debug('Formatting output: %s', output) - - if not output: - raise GenericWarning('No output') - - # Ensure we're working with a list - if isinstance(output, subprocess.CompletedProcess): - stdout = output.stdout - if not isinstance(stdout, str): - stdout = stdout.decode('utf8') - output = stdout.strip().splitlines() - if not output: - # Going to treat these as successes (for now) - LOG.warning('Program output was empty, assuming good result.') - return color_string(msg_good, 'GREEN') - else: - try: - output = list(output) - except TypeError: - output = [output] - - # Safety check - if not output: - # Going to ignore empty function output for now - LOG.error('Output is empty') - return 'UNKNOWN' - - # Build result_msg - result_msg = f'{output.pop(0)}' - if output: - output = [f'{" "*(self.indent+self.width)}{line}' for line in output] - result_msg += '\n' + '\n'.join(output) - - # Done - return result_msg - - def _log_result(self, message, result_msg): - """Log result text without color formatting.""" - log_text = f'{" "*self.indent}{message:<{self.width}}{result_msg}' - for line in log_text.splitlines(): - line = strip_colors(line) - LOG.info(line) - - def add_error(self, exception_name): - """Add exception name to error list.""" - if exception_name not in self.list_errors: - self.list_errors.append(exception_name) - - def add_warning(self, exception_name): - """Add exception name to warning list.""" - if exception_name not in self.list_warnings: - self.list_warnings.append(exception_name) - - def run( - self, message, function, *args, - catch_all=None, msg_good=None, verbose=None, **kwargs): - """Run a function and print the results, returns results as dict. - - If catch_all is True then (nearly) all exceptions will be caught. - Otherwise if an exception occurs that wasn't specified it will be - re-raised. - - If the function returns data it will be used instead of msg_good, - msg_bad, or exception text. - The output should be a list or a subprocess.CompletedProcess object. - - If msg_good is passed it will override self.msg_good for this call. - - If verbose is True then exception names or messages will be used for - the result message. Otherwise it will simply be set to result_bad. - - If catch_all and/or verbose are passed it will override - self.catch_all and/or self.verbose for this call. - - args and kwargs are passed to the function. - """ - LOG.debug('function: %s.%s', function.__module__, function.__name__) - LOG.debug('args: %s', args) - LOG.debug('kwargs: %s', kwargs) - LOG.debug( - 'catch_all: %s, msg_good: %s, verbose: %s', - catch_all, - msg_good, - verbose, - ) - f_exception = None - catch_all = catch_all if catch_all is not None else self.catch_all - msg_good = msg_good if msg_good is not None else self.msg_good - output = None - result_msg = 'UNKNOWN' - verbose = verbose if verbose is not None else self.verbose - - # Build exception tuples - e_exceptions = tuple(get_exception(e) for e in self.list_errors) - w_exceptions = tuple(get_exception(e) for e in self.list_warnings) - - # Run function and catch exceptions - print(f'{" "*self.indent}{message:<{self.width}}', end='', flush=True) - LOG.debug('Running function: %s.%s', function.__module__, function.__name__) - try: - output = function(*args, **kwargs) - except w_exceptions as _exception: - # Warnings - result_msg = self._format_exception_message(_exception) - print_warning(result_msg, log=False) - f_exception = _exception - except e_exceptions as _exception: - # Exceptions - result_msg = self._format_exception_message(_exception) - print_error(result_msg, log=False) - f_exception = _exception - except Exception as _exception: - # Unexpected exceptions - if verbose: - result_msg = self._format_exception_message(_exception) - else: - result_msg = self.msg_bad - print_error(result_msg, log=False) - f_exception = _exception - if not catch_all: - # Re-raise error as necessary - raise - else: - # Success - if output: - result_msg = self._format_function_output(output, msg_good) - print(result_msg) - else: - result_msg = msg_good - print_success(result_msg, log=False) - - # Done - self._log_result(message, result_msg) - return { - 'Exception': f_exception, - 'Failed': bool(f_exception), - 'Message': result_msg, - 'Output': output, - } - - # Functions -def abort(prompt='Aborted.', show_prompt=True, return_code=1): - """Abort script.""" - print_warning(prompt) - if show_prompt: - sleep(0.5) - pause(prompt='Press Enter to exit... ') - sys.exit(return_code) - - -def ask(prompt='Kotaero!'): - """Prompt the user with a Y/N question, returns bool.""" - answer = None - prompt = f'{prompt} [Y/N]: ' - - # Loop until acceptable answer is given - while answer is None: - tmp = input_text(prompt) - if re.search(r'^y(es|up|)$', tmp, re.IGNORECASE): - answer = True - elif re.search(r'^n(o|ope|)$', tmp, re.IGNORECASE): - answer = False - - # Done - LOG.info('%s%s', prompt, 'Yes' if answer else 'No') - return answer - - -def beep(repeat=1): - """Play system bell with optional repeat.""" - while repeat >= 1: - # Print bell char without a newline - print('\a', end='', flush=True) - sleep(0.5) - repeat -= 1 - - def bytes_to_string(size, decimals=0, use_binary=True): """Convert size into a human-readable format, returns str. @@ -682,323 +73,6 @@ def bytes_to_string(size, decimals=0, use_binary=True): return size_str -def choice(choices, prompt='答えろ！'): - """Choose an option from a provided list, returns str. - - Choices provided will be converted to uppercase and returned as such. - Similar to the commands choice (Windows) and select (Linux). - """ - LOG.debug('choices: %s, prompt: %s', choices, prompt) - answer = None - choices = [str(c).upper()[:1] for c in choices] - prompt = f'{prompt} [{"/".join(choices)}]' - regex = f'^({"|".join(choices)})$' - - # Loop until acceptable answer is given - while answer is None: - tmp = input_text(prompt=prompt) - if re.search(regex, tmp, re.IGNORECASE): - answer = tmp.upper() - - # Done - LOG.info('%s %s', prompt, answer) - return answer - - -def clear_screen(): - """Simple wrapper for clear/cls.""" - cmd = 'cls' if os.name == 'nt' else 'clear' - proc = subprocess.run(cmd, check=False, shell=True, stderr=subprocess.PIPE) - - # Workaround for live macOS env - if proc.returncode != 0: - print('\033c') - - -def color_string(strings, colors, sep=' '): - """Build colored string using ANSI escapes, returns str.""" - clear_code = COLORS['CLEAR'] - msg = [] - - # Convert to tuples if necessary - if isinstance(strings, (str, pathlib.Path)): - strings = (strings,) - if isinstance(colors, (str, pathlib.Path)): - colors = (colors,) - - # Convert to strings if necessary - try: - iter(strings) - except TypeError: - # Assuming single element passed, convert to string - strings = (str(strings),) - try: - iter(colors) - except TypeError: - # Assuming single element passed, convert to string - colors = (str(colors),) - - # Build new string with color escapes added - for string, color in itertools.zip_longest(strings, colors): - color_code = COLORS.get(color, clear_code) - msg.append(f'{color_code}{string}{clear_code}') - - # Done - return sep.join(msg) - - -def generate_debug_report(): - """Generate debug report, returns str.""" - platform_function_list = ( - 'architecture', - 'machine', - 'platform', - 'python_version', - ) - report = [] - - # Logging data - log_path = get_log_filepath() - if log_path: - report.append('------ Start Log -------') - report.append('') - with open(log_path, 'r', encoding='utf-8') as log_file: - report.extend(log_file.read().splitlines()) - report.append('') - report.append('------- End Log --------') - - # System - report.append('--- Start debug info ---') - report.append('') - report.append('[System]') - report.append(f' {"FQDN":<24} {socket.getfqdn()}') - for func in platform_function_list: - func_name = func.replace('_', ' ').capitalize() - func_result = getattr(platform, func)() - report.append(f' {func_name:<24} {func_result}') - report.append(f' {"Python sys.argv":<24} {sys.argv}') - report.append('') - - # Environment - report.append('[Environment Variables]') - for key, value in sorted(os.environ.items()): - report.append(f' {key:<24} {value}') - report.append('') - - # Done - report.append('---- End debug info ----') - return '\n'.join(report) - - -@cache -def get_exception(name): - """Get exception by name, returns exception object. - - [Doctest] - >>> t = TryAndPrint() - >>> t._get_exception('AttributeError') - - >>> t._get_exception('CalledProcessError') - - >>> t._get_exception('GenericError') - - """ - LOG.debug('Getting exception: %s', name) - obj = getattr(sys.modules[__name__], name, None) - if obj: - return obj - - # Try builtin classes - obj = getattr(sys.modules['builtins'], name, None) - if obj: - return obj - - # Try all modules - for _mod in sys.modules.values(): - obj = getattr(_mod, name, None) - if obj: - break - - # Check if not found - if not obj: - raise AttributeError(f'Failed to find exception: {name}') - - # Done - return obj - - -def get_log_filepath(): - """Get the log filepath from the root logger, returns pathlib.Path obj. - - NOTE: This will use the first handler baseFilename it finds (if any). - """ - log_filepath = None - root_logger = logging.getLogger() - - # Check handlers - for handler in root_logger.handlers: - if hasattr(handler, 'baseFilename'): - log_filepath = pathlib.Path(handler.baseFilename).resolve() - break - - # Done - return log_filepath - - -def input_text(prompt='Enter text', allow_empty_response=True): - """Get text from user, returns string.""" - prompt = str(prompt) - response = None - if prompt[-1:] != ' ': - prompt += ' ' - print(prompt, end='', flush=True) - - while response is None: - try: - response = input() - LOG.debug('%s%s', prompt, response) - except EOFError: - # Ignore and try again - LOG.warning('Exception occured', exc_info=True) - print('', flush=True) - if not allow_empty_response: - if response is None or not response.strip(): - # The None check here is used to avoid a TypeError if response is None - print(f'\r{prompt}', end='', flush=True) - response = None - - return response - - -def major_exception(): - """Display traceback, optionally upload detailes, and exit.""" - LOG.critical('Major exception encountered', exc_info=True) - print_error('Major exception', log=False) - print_warning(SUPPORT_MESSAGE) - if ENABLED_UPLOAD_DATA: - print_warning('Also, please run upload-logs to help debugging!') - print(traceback.format_exc()) - - # Done - pause('Press Enter to exit... ') - raise SystemExit(1) - - -def pause(prompt='Press Enter to continue... '): - """Simple pause implementation.""" - input_text(prompt) - - -def print_colored(strings, colors, log=False, sep=' ', **kwargs): - """Prints strings in the colors specified.""" - LOG.debug( - 'strings: %s, colors: %s, sep: %s, kwargs: %s', - strings, colors, sep, kwargs, - ) - msg = color_string(strings, colors, sep=sep) - print_options = { - 'end': kwargs.get('end', '\n'), - 'file': kwargs.get('file', sys.stdout), - 'flush': kwargs.get('flush', False), - } - - print(msg, **print_options) - if log: - LOG.info(strip_colors(msg)) - - -def print_error(msg, log=True, **kwargs): - """Prints message in RED and log as ERROR.""" - if 'file' not in kwargs: - # Only set if not specified - kwargs['file'] = sys.stderr - print_colored(msg, 'RED', **kwargs) - if log: - LOG.error(msg) - - -def print_info(msg, log=True, **kwargs): - """Prints message in BLUE and log as INFO.""" - print_colored(msg, 'BLUE', **kwargs) - if log: - LOG.info(msg) - - -def print_report(report, indent=None, log=True): - """Print report to screen and optionally to log.""" - for line in report: - if indent: - line = f'{" "*indent}{line}' - print(line) - if log: - LOG.info(strip_colors(line)) - - -def print_standard(msg, log=True, **kwargs): - """Prints message and log as INFO.""" - print(msg, **kwargs) - if log: - LOG.info(msg) - - -def print_success(msg, log=True, **kwargs): - """Prints message in GREEN and log as INFO.""" - print_colored(msg, 'GREEN', **kwargs) - if log: - LOG.info(msg) - - -def print_warning(msg, log=True, **kwargs): - """Prints message in YELLOW and log as WARNING.""" - if 'file' not in kwargs: - # Only set if not specified - kwargs['file'] = sys.stderr - print_colored(msg, 'YELLOW', **kwargs) - if log: - LOG.warning(msg) - - -def save_pickles(obj_dict, out_path=None): - """Save dict of objects using pickle.""" - LOG.info('Saving pickles') - - # Set path - if not out_path: - out_path = pathlib.Path(f'{get_root_logger_path().parent}/debug') - - # Save pickles - try: - for name, obj in obj_dict.copy().items(): - if name.startswith('__') or inspect.ismodule(obj): - continue - with open(f'{out_path}/{name}.pickle', 'wb') as _f: - pickle.dump(obj, _f, protocol=pickle.HIGHEST_PROTOCOL) - except Exception: - LOG.error('Failed to save all the pickles', exc_info=True) - - -def set_title(title): - """Set window title.""" - LOG.debug('title: %s', title) - if os.name == 'nt': - os.system(f'title {title}') - else: - print_error('Setting the title is only supported under Windows.') - - -def show_data(message, data, color=None, indent=None, width=None): - """Display info using default or provided indent and width.""" - colors = (None, color if color else None) - indent = INDENT if indent is None else indent - width = WIDTH if width is None else width - print_colored( - (f'{" "*indent}{message:<{width}}', data), - colors, - log=True, - sep='', - ) - - def sleep(seconds=2): """Simple wrapper for time.sleep.""" time.sleep(seconds) @@ -1053,55 +127,5 @@ def string_to_bytes(size, assume_binary=False): return size -def strip_colors(string): - """Strip known ANSI color escapes from string, returns str.""" - LOG.debug('string: %s', string) - for color in COLORS.values(): - string = string.replace(color, '') - return string - - -def upload_debug_report(report, compress=True, reason='DEBUG'): - """Upload debug report to CRASH_SERVER as specified in wk.cfg.main.""" - LOG.info('Uploading debug report to %s', CRASH_SERVER.get('Name', '?')) - headers = CRASH_SERVER.get('Headers', {'X-Requested-With': 'XMLHttpRequest'}) - if compress: - headers['Content-Type'] = 'application/octet-stream' - - # Check if the required server details are available - if not all(CRASH_SERVER.get(key, False) for key in ('Name', 'Url', 'User')): - msg = 'Server details missing, aborting upload.' - print_error(msg) - raise RuntimeError(msg) - - # Set filename (based on the logging config if possible) - filename = 'Unknown' - log_path = get_log_filepath() - if log_path: - # Strip everything but the prefix - filename = re.sub(r'^(.*)_(\d{4}-\d{2}-\d{2}.*)', r'\1', log_path.name) - filename = f'{filename}_{reason}_{time.strftime("%Y-%m-%d_%H%M%S%z")}.log' - LOG.debug('filename: %s', filename) - - # Compress report - if compress: - filename += '.xz' - xz_report = lzma.compress(report.encode('utf8')) - - # Upload report - url = f'{CRASH_SERVER["Url"]}/{filename}' - response = requests.put( - url, - auth=(CRASH_SERVER['User'], CRASH_SERVER.get('Pass', '')), - data=xz_report if compress else report, - headers=headers, - timeout=60, - ) - - # Check response - if not response.ok: - raise RuntimeError('Failed to upload report') - - if __name__ == '__main__': print("This file is not meant to be called directly.") diff --git a/scripts/wk/ui/__init__.py b/scripts/wk/ui/__init__.py new file mode 100644 index 00000000..eb6fa151 --- /dev/null +++ b/scripts/wk/ui/__init__.py @@ -0,0 +1,3 @@ +"""WizardKit: ui module init""" + +from . import cli diff --git a/scripts/wk/ui/cli.py b/scripts/wk/ui/cli.py new file mode 100644 index 00000000..ecd8965a --- /dev/null +++ b/scripts/wk/ui/cli.py @@ -0,0 +1,862 @@ +"""WizardKit: CLI functions""" +# vim: sts=2 sw=2 ts=2 + +import itertools +import logging +import os +import pathlib +import platform +import re +import subprocess +import sys +import traceback + +from collections import OrderedDict + +try: + from functools import cache +except ImportError: + # Assuming Python is < 3.9 + from functools import lru_cache as cache + +from wk.cfg.main import ( + ENABLED_UPLOAD_DATA, + INDENT, + SUPPORT_MESSAGE, + WIDTH, + ) +from wk.std import (sleep, GenericWarning) + +# STATIC VARIABLES +COLORS = { + 'CLEAR': '\033[0m', + 'RED': '\033[31m', + 'RED_BLINK': '\033[31;5m', + 'ORANGE': '\033[31;1m', + 'ORANGE_RED': '\033[1;31;41m', + 'GREEN': '\033[32m', + 'YELLOW': '\033[33m', + 'YELLOW_BLINK': '\033[33;5m', + 'BLUE': '\033[34m', + 'PURPLE': '\033[35m', + 'CYAN': '\033[36m', + } +LOG = logging.getLogger(__name__) +PLATFORM = platform.system() + + +# Classes +class Menu(): + """Object for tracking menu specific data and methods. + + Menu items are added to an OrderedDict so the order is preserved. + + ASSUMPTIONS: + 1. All entry names are unique. + 2. All action entry names start with different letters. + """ + def __init__(self, title='[Untitled Menu]'): + self.actions = OrderedDict() + self.options = OrderedDict() + self.sets = OrderedDict() + self.toggles = OrderedDict() + self.disabled_str = 'Disabled' + self.separator = '─' + self.title = title + + def _generate_menu_text(self): + """Generate menu text, returns str.""" + separator_string = self._get_separator_string() + menu_lines = [self.title, separator_string] if self.title else [] + + # Sets & toggles + for section in (self.sets, self.toggles): + for details in section.values(): + if details.get('Hidden', False): + continue + if details.get('Separator', False): + menu_lines.append(separator_string) + menu_lines.append(details['Display Name']) + if self.sets or self.toggles: + menu_lines.append(separator_string) + + # Options + for details in self.options.values(): + if details.get('Hidden', False): + continue + if details.get('Separator', False): + menu_lines.append(separator_string) + menu_lines.append(details['Display Name']) + if self.options: + menu_lines.append(separator_string) + + # Actions + for details in self.actions.values(): + if details.get('Hidden', False): + continue + if details.get('Separator', False): + menu_lines.append(separator_string) + menu_lines.append(details['Display Name']) + + # Show menu + menu_lines.append('') + menu_lines = [str(line) for line in menu_lines] + return '\n'.join(menu_lines) + + def _get_display_name( + self, name, details, + index=None, no_checkboxes=True, setting_item=False): + """Format display name based on details and args, returns str.""" + disabled = details.get('Disabled', False) + if setting_item and not details['Selected']: + # Display item in YELLOW + disabled = True + checkmark = '*' + if 'DISPLAY' in os.environ or PLATFORM == 'Darwin': + checkmark = '✓' + display_name = f'{index if index else name[:1].upper()}: ' + if not (index and index >= 10): + display_name = f' {display_name}' + if setting_item and 'Value' in details: + name = f'{name} = {details["Value"]}' + + # Add enabled status if necessary + if not no_checkboxes: + display_name += f'[{checkmark if details["Selected"] else " "}] ' + + # Add name + if disabled: + display_name += color_string(f'{name} ({self.disabled_str})', 'YELLOW') + else: + display_name += name + + # Done + return display_name + + def _get_separator_string(self): + """Format separator length based on name lengths, returns str.""" + separator_length = 0 + + # Check title line(s) + if self.title: + for line in self.title.split('\n'): + separator_length = max(separator_length, len(strip_colors(line))) + + # Loop over all item names + for section in (self.actions, self.options, self.sets, self.toggles): + for details in section.values(): + if details.get('Hidden', False): + # Skip hidden lines + continue + line = strip_colors(details['Display Name']) + separator_length = max(separator_length, len(line)) + separator_length += 1 + + # Done + return self.separator * separator_length + + def _get_valid_answers(self): + """Get valid answers based on menu items, returns list.""" + valid_answers = [] + + # Numbered items + index = 0 + for section in (self.sets, self.toggles, self.options): + for details in section.values(): + if details.get('Hidden', False): + # Don't increment index or add to valid_answers + continue + index += 1 + if not details.get('Disabled', False): + valid_answers.append(str(index)) + + # Action items + for name, details in self.actions.items(): + if not details.get('Disabled', False): + valid_answers.append(name[:1].upper()) + + # Done + return valid_answers + + def _resolve_selection(self, selection): + """Get menu item based on user selection, returns tuple.""" + offset = 1 + resolved_selection = None + if selection.isnumeric(): + # Enumerate over numbered entries + entries = [ + *self.sets.items(), + *self.toggles.items(), + *self.options.items(), + ] + for _i, details in enumerate(entries): + if details[1].get('Hidden', False): + offset -= 1 + elif str(_i+offset) == selection: + resolved_selection = (details) + break + else: + # Just check actions + for action, details in self.actions.items(): + if action.lower().startswith(selection.lower()): + resolved_selection = (action, details) + break + + # Done + return resolved_selection + + def _update(self, single_selection=True, settings_mode=False): + """Update menu items in preparation for printing to screen.""" + index = 0 + + # Fix selection status for sets + for set_details in self.sets.values(): + set_selected = True + set_targets = set_details['Targets'] + for option, option_details in self.options.items(): + if option in set_targets and not option_details['Selected']: + set_selected = False + elif option not in set_targets and option_details['Selected']: + set_selected = False + set_details['Selected'] = set_selected + + # Numbered sections + for section in (self.sets, self.toggles, self.options): + for name, details in section.items(): + if details.get('Hidden', False): + # Skip hidden lines and don't increment index + continue + index += 1 + details['Display Name'] = self._get_display_name( + name, + details, + index=index, + no_checkboxes=single_selection, + setting_item=settings_mode, + ) + + # Actions + for name, details in self.actions.items(): + details['Display Name'] = self._get_display_name( + name, + details, + no_checkboxes=True, + ) + + def _update_entry_selection_status(self, entry, toggle=True, status=None): + """Update entry selection status either directly or by toggling.""" + if entry in self.sets: + # Update targets not the set itself + new_status = not self.sets[entry]['Selected'] if toggle else status + targets = self.sets[entry]['Targets'] + self._update_set_selection_status(targets, new_status) + for section in (self.toggles, self.options, self.actions): + if entry in section: + if toggle: + section[entry]['Selected'] = not section[entry]['Selected'] + else: + section[entry]['Selected'] = status + + def _update_set_selection_status(self, targets, status): + """Select or deselect options based on targets and status.""" + for option, details in self.options.items(): + # If (new) status is True and this option is a target then select + # Otherwise deselect + details['Selected'] = status and option in targets + + def _user_select(self, prompt): + """Show menu and select an entry, returns str.""" + menu_text = self._generate_menu_text() + valid_answers = self._get_valid_answers() + + # Menu loop + while True: + clear_screen() + print(menu_text) + sleep(0.01) + answer = input_text(prompt).strip() + if answer.upper() in valid_answers: + break + + # Done + return answer + + def add_action(self, name, details=None): + """Add action to menu.""" + details = details if details else {} + details['Selected'] = details.get('Selected', False) + self.actions[name] = details + + def add_option(self, name, details=None): + """Add option to menu.""" + details = details if details else {} + details['Selected'] = details.get('Selected', False) + self.options[name] = details + + def add_set(self, name, details=None): + """Add set to menu.""" + details = details if details else {} + details['Selected'] = details.get('Selected', False) + + # Safety check + if 'Targets' not in details: + raise KeyError('Menu set has no targets') + + # Add set + self.sets[name] = details + + def add_toggle(self, name, details=None): + """Add toggle to menu.""" + details = details if details else {} + details['Selected'] = details.get('Selected', False) + self.toggles[name] = details + + def advanced_select(self, prompt='Please make a selection: '): + """Display menu and make multiple selections, returns tuple. + + NOTE: Menu is displayed until an action entry is selected. + """ + while True: + self._update(single_selection=False) + user_selection = self._user_select(prompt) + selected_entry = self._resolve_selection(user_selection) + if user_selection.isnumeric(): + # Update selection(s) + self._update_entry_selection_status(selected_entry[0]) + else: + # Action selected + break + + # Done + return selected_entry + + def settings_select(self, prompt='Please make a selection: '): + """Display menu and make multiple selections, returns tuple. + + NOTE: Menu is displayed until an action entry is selected. + """ + choice_kwargs = { + 'choices': ['T', 'C'], + 'prompt': 'Toggle or change value?', + } + + while True: + self._update(single_selection=True, settings_mode=True) + user_selection = self._user_select(prompt) + selected_entry = self._resolve_selection(user_selection) + if user_selection.isnumeric(): + if 'Value' in selected_entry[-1] and choice(**choice_kwargs) == 'C': + # Change + selected_entry[-1]['Value'] = input_text('Enter new value: ') + else: + # Toggle + self._update_entry_selection_status(selected_entry[0]) + else: + # Action selected + break + + # Done + return selected_entry + + def simple_select(self, prompt='Please make a selection: ', update=True): + """Display menu and make a single selection, returns tuple.""" + if update: + self._update() + user_selection = self._user_select(prompt) + return self._resolve_selection(user_selection) + + def update(self): + """Update menu with default settings.""" + self._update() + + +class TryAndPrint(): + """Object used to standardize running functions and returning the result. + + The errors and warning attributes are used to allow fine-tuned results + based on exception names. + """ + def __init__(self, msg_bad='FAILED', msg_good='SUCCESS'): + self.catch_all = True + self.indent = INDENT + self.list_errors = ['GenericError'] + self.list_warnings = ['GenericWarning'] + self.msg_bad = msg_bad + self.msg_good = msg_good + self.verbose = False + self.width = WIDTH + + def _format_exception_message(self, _exception): + """Format using the exception's args or name, returns str.""" + LOG.debug( + 'Formatting exception: %s, %s', + _exception.__class__.__name__, + _exception, + ) + message = '' + + # Format message string from _exception + try: + if isinstance(_exception, subprocess.CalledProcessError): + message = _exception.stderr + if not isinstance(message, str): + message = message.decode('utf-8') + message = message.strip() + elif isinstance(_exception, ZeroDivisionError): + # Skip and just use exception name below + pass + else: + message = str(_exception) + except Exception: + # Just use the exception name instead + pass + + # Prepend exception name + if _exception.__class__.__name__ not in ('GenericError', 'GenericWarning'): + try: + message = f'{_exception.__class__.__name__}: {message}' + except Exception: + message = f'UNKNOWN ERROR: {message}' + + # Fix multi-line messages + if '\n' in message: + try: + lines = [ + f'{" "*(self.indent+self.width)}{line.strip()}' + for line in message.splitlines() if line.strip() + ] + lines[0] = lines[0].strip() + message = '\n'.join(lines) + except Exception: + pass + + # Done + return message + + def _format_function_output(self, output, msg_good): + """Format function output for use in try_and_print(), returns str.""" + LOG.debug('Formatting output: %s', output) + + if not output: + raise GenericWarning('No output') + + # Ensure we're working with a list + if isinstance(output, subprocess.CompletedProcess): + stdout = output.stdout + if not isinstance(stdout, str): + stdout = stdout.decode('utf8') + output = stdout.strip().splitlines() + if not output: + # Going to treat these as successes (for now) + LOG.warning('Program output was empty, assuming good result.') + return color_string(msg_good, 'GREEN') + else: + try: + output = list(output) + except TypeError: + output = [output] + + # Safety check + if not output: + # Going to ignore empty function output for now + LOG.error('Output is empty') + return 'UNKNOWN' + + # Build result_msg + result_msg = f'{output.pop(0)}' + if output: + output = [f'{" "*(self.indent+self.width)}{line}' for line in output] + result_msg += '\n' + '\n'.join(output) + + # Done + return result_msg + + def _log_result(self, message, result_msg): + """Log result text without color formatting.""" + log_text = f'{" "*self.indent}{message:<{self.width}}{result_msg}' + for line in log_text.splitlines(): + line = strip_colors(line) + LOG.info(line) + + def add_error(self, exception_name): + """Add exception name to error list.""" + if exception_name not in self.list_errors: + self.list_errors.append(exception_name) + + def add_warning(self, exception_name): + """Add exception name to warning list.""" + if exception_name not in self.list_warnings: + self.list_warnings.append(exception_name) + + def run( + self, message, function, *args, + catch_all=None, msg_good=None, verbose=None, **kwargs): + """Run a function and print the results, returns results as dict. + + If catch_all is True then (nearly) all exceptions will be caught. + Otherwise if an exception occurs that wasn't specified it will be + re-raised. + + If the function returns data it will be used instead of msg_good, + msg_bad, or exception text. + The output should be a list or a subprocess.CompletedProcess object. + + If msg_good is passed it will override self.msg_good for this call. + + If verbose is True then exception names or messages will be used for + the result message. Otherwise it will simply be set to result_bad. + + If catch_all and/or verbose are passed it will override + self.catch_all and/or self.verbose for this call. + + args and kwargs are passed to the function. + """ + LOG.debug('function: %s.%s', function.__module__, function.__name__) + LOG.debug('args: %s', args) + LOG.debug('kwargs: %s', kwargs) + LOG.debug( + 'catch_all: %s, msg_good: %s, verbose: %s', + catch_all, + msg_good, + verbose, + ) + f_exception = None + catch_all = catch_all if catch_all is not None else self.catch_all + msg_good = msg_good if msg_good is not None else self.msg_good + output = None + result_msg = 'UNKNOWN' + verbose = verbose if verbose is not None else self.verbose + + # Build exception tuples + e_exceptions = tuple(get_exception(e) for e in self.list_errors) + w_exceptions = tuple(get_exception(e) for e in self.list_warnings) + + # Run function and catch exceptions + print(f'{" "*self.indent}{message:<{self.width}}', end='', flush=True) + LOG.debug('Running function: %s.%s', function.__module__, function.__name__) + try: + output = function(*args, **kwargs) + except w_exceptions as _exception: + # Warnings + result_msg = self._format_exception_message(_exception) + print_warning(result_msg, log=False) + f_exception = _exception + except e_exceptions as _exception: + # Exceptions + result_msg = self._format_exception_message(_exception) + print_error(result_msg, log=False) + f_exception = _exception + except Exception as _exception: + # Unexpected exceptions + if verbose: + result_msg = self._format_exception_message(_exception) + else: + result_msg = self.msg_bad + print_error(result_msg, log=False) + f_exception = _exception + if not catch_all: + # Re-raise error as necessary + raise + else: + # Success + if output: + result_msg = self._format_function_output(output, msg_good) + print(result_msg) + else: + result_msg = msg_good + print_success(result_msg, log=False) + + # Done + self._log_result(message, result_msg) + return { + 'Exception': f_exception, + 'Failed': bool(f_exception), + 'Message': result_msg, + 'Output': output, + } + + +# Functions +def abort(prompt='Aborted.', show_prompt=True, return_code=1): + """Abort script.""" + print_warning(prompt) + if show_prompt: + sleep(0.5) + pause(prompt='Press Enter to exit... ') + sys.exit(return_code) + + +def ask(prompt='Kotaero!'): + """Prompt the user with a Y/N question, returns bool.""" + answer = None + prompt = f'{prompt} [Y/N]: ' + + # Loop until acceptable answer is given + while answer is None: + tmp = input_text(prompt) + if re.search(r'^y(es|up|)$', tmp, re.IGNORECASE): + answer = True + elif re.search(r'^n(o|ope|)$', tmp, re.IGNORECASE): + answer = False + + # Done + LOG.info('%s%s', prompt, 'Yes' if answer else 'No') + return answer + + +def beep(repeat=1): + """Play system bell with optional repeat.""" + while repeat >= 1: + # Print bell char without a newline + print('\a', end='', flush=True) + sleep(0.5) + repeat -= 1 + + +def choice(choices, prompt='答えろ！'): + """Choose an option from a provided list, returns str. + + Choices provided will be converted to uppercase and returned as such. + Similar to the commands choice (Windows) and select (Linux). + """ + LOG.debug('choices: %s, prompt: %s', choices, prompt) + answer = None + choices = [str(c).upper()[:1] for c in choices] + prompt = f'{prompt} [{"/".join(choices)}]' + regex = f'^({"|".join(choices)})$' + + # Loop until acceptable answer is given + while answer is None: + tmp = input_text(prompt=prompt) + if re.search(regex, tmp, re.IGNORECASE): + answer = tmp.upper() + + # Done + LOG.info('%s %s', prompt, answer) + return answer + + +def clear_screen(): + """Simple wrapper for clear/cls.""" + cmd = 'cls' if os.name == 'nt' else 'clear' + proc = subprocess.run(cmd, check=False, shell=True, stderr=subprocess.PIPE) + + # Workaround for live macOS env + if proc.returncode != 0: + print('\033c') + + +def color_string(strings, colors, sep=' '): + """Build colored string using ANSI escapes, returns str.""" + clear_code = COLORS['CLEAR'] + msg = [] + + # Convert to tuples if necessary + if isinstance(strings, (str, pathlib.Path)): + strings = (strings,) + if isinstance(colors, (str, pathlib.Path)): + colors = (colors,) + + # Convert to strings if necessary + try: + iter(strings) + except TypeError: + # Assuming single element passed, convert to string + strings = (str(strings),) + try: + iter(colors) + except TypeError: + # Assuming single element passed, convert to string + colors = (str(colors),) + + # Build new string with color escapes added + for string, color in itertools.zip_longest(strings, colors): + color_code = COLORS.get(color, clear_code) + msg.append(f'{color_code}{string}{clear_code}') + + # Done + return sep.join(msg) + + +@cache +def get_exception(name): + """Get exception by name, returns exception object. + + [Doctest] + >>> t = TryAndPrint() + >>> t._get_exception('AttributeError') + + >>> t._get_exception('CalledProcessError') + + >>> t._get_exception('GenericError') + + """ + LOG.debug('Getting exception: %s', name) + obj = getattr(sys.modules[__name__], name, None) + if obj: + return obj + + # Try builtin classes + obj = getattr(sys.modules['builtins'], name, None) + if obj: + return obj + + # Try all modules + for _mod in sys.modules.values(): + obj = getattr(_mod, name, None) + if obj: + break + + # Check if not found + if not obj: + raise AttributeError(f'Failed to find exception: {name}') + + # Done + return obj + + +def input_text(prompt='Enter text', allow_empty_response=True): + """Get text from user, returns string.""" + prompt = str(prompt) + response = None + if prompt[-1:] != ' ': + prompt += ' ' + print(prompt, end='', flush=True) + + while response is None: + try: + response = input() + LOG.debug('%s%s', prompt, response) + except EOFError: + # Ignore and try again + LOG.warning('Exception occured', exc_info=True) + print('', flush=True) + if not allow_empty_response: + if response is None or not response.strip(): + # The None check here is used to avoid a TypeError if response is None + print(f'\r{prompt}', end='', flush=True) + response = None + + return response + + +def major_exception(): + """Display traceback, optionally upload detailes, and exit.""" + LOG.critical('Major exception encountered', exc_info=True) + print_error('Major exception', log=False) + print_warning(SUPPORT_MESSAGE) + if ENABLED_UPLOAD_DATA: + print_warning('Also, please run upload-logs to help debugging!') + print(traceback.format_exc()) + + # Done + pause('Press Enter to exit... ') + raise SystemExit(1) + + +def pause(prompt='Press Enter to continue... '): + """Simple pause implementation.""" + input_text(prompt) + + +def print_colored(strings, colors, log=False, sep=' ', **kwargs): + """Prints strings in the colors specified.""" + LOG.debug( + 'strings: %s, colors: %s, sep: %s, kwargs: %s', + strings, colors, sep, kwargs, + ) + msg = color_string(strings, colors, sep=sep) + print_options = { + 'end': kwargs.get('end', '\n'), + 'file': kwargs.get('file', sys.stdout), + 'flush': kwargs.get('flush', False), + } + + print(msg, **print_options) + if log: + LOG.info(strip_colors(msg)) + + +def print_error(msg, log=True, **kwargs): + """Prints message in RED and log as ERROR.""" + if 'file' not in kwargs: + # Only set if not specified + kwargs['file'] = sys.stderr + print_colored(msg, 'RED', **kwargs) + if log: + LOG.error(msg) + + +def print_info(msg, log=True, **kwargs): + """Prints message in BLUE and log as INFO.""" + print_colored(msg, 'BLUE', **kwargs) + if log: + LOG.info(msg) + + +def print_report(report, indent=None, log=True): + """Print report to screen and optionally to log.""" + for line in report: + if indent: + line = f'{" "*indent}{line}' + print(line) + if log: + LOG.info(strip_colors(line)) + + +def print_standard(msg, log=True, **kwargs): + """Prints message and log as INFO.""" + print(msg, **kwargs) + if log: + LOG.info(msg) + + +def print_success(msg, log=True, **kwargs): + """Prints message in GREEN and log as INFO.""" + print_colored(msg, 'GREEN', **kwargs) + if log: + LOG.info(msg) + + +def print_warning(msg, log=True, **kwargs): + """Prints message in YELLOW and log as WARNING.""" + if 'file' not in kwargs: + # Only set if not specified + kwargs['file'] = sys.stderr + print_colored(msg, 'YELLOW', **kwargs) + if log: + LOG.warning(msg) + + +def set_title(title): + """Set window title.""" + LOG.debug('title: %s', title) + if os.name == 'nt': + os.system(f'title {title}') + else: + print_error('Setting the title is only supported under Windows.') + + +def show_data(message, data, color=None, indent=None, width=None): + """Display info using default or provided indent and width.""" + colors = (None, color if color else None) + indent = INDENT if indent is None else indent + width = WIDTH if width is None else width + print_colored( + (f'{" "*indent}{message:<{width}}', data), + colors, + log=True, + sep='', + ) + + +def strip_colors(string): + """Strip known ANSI color escapes from string, returns str.""" + LOG.debug('string: %s', string) + for color in COLORS.values(): + string = string.replace(color, '') + return string + + +if __name__ == '__main__': + print("This file is not meant to be called directly.") diff --git a/scripts/wk_debug.py b/scripts/wk_debug.py index 86900004..1808be89 100755 --- a/scripts/wk_debug.py +++ b/scripts/wk_debug.py @@ -17,7 +17,7 @@ OPTIONS = { def get_debug_prefix() -> str: """Ask what we're debugging, returns log dir prefix.""" - menu = wk.std.Menu(title=f'{wk.cfg.main.KIT_NAME_FULL}: Debugging Menu\n') + menu = wk.ui.cli.Menu(title=f'{wk.cfg.main.KIT_NAME_FULL}: Debugging Menu\n') for name, prefix in OPTIONS.items(): menu.add_option(name, {'Prefix': prefix}) selection = menu.simple_select() @@ -38,14 +38,14 @@ def get_debug_path() -> pathlib.Path: # Safety check if not debug_paths: - wk.std.abort('No logs found, aborting.') + wk.ui.cli.abort('No logs found, aborting.') # Use latest option - if wk.std.ask('Use latest session?'): + if wk.ui.cli.ask('Use latest session?'): return debug_paths[-1] # Select from list - menu = wk.std.Menu(title=f'{wk.cfg.main.KIT_NAME_FULL}: Debugging Menu\n') + menu = wk.ui.cli.Menu(title=f'{wk.cfg.main.KIT_NAME_FULL}: Debugging Menu\n') for item in debug_paths: menu.add_option(item.parent.name, {'Path': item}) selection = menu.simple_select() From 6efc970374235c48526e0df4a9f1799f42c2d0c3 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sat, 1 Apr 2023 22:36:59 -0700 Subject: [PATCH 02/16] Move tmux sections into wk.ui --- scripts/wk/__init__.py | 1 - scripts/wk/clone/ddrescue.py | 3 ++- scripts/wk/hw/cpu.py | 2 +- scripts/wk/hw/diags.py | 3 ++- scripts/wk/hw/screensavers.py | 6 +++--- scripts/wk/ui/__init__.py | 1 + scripts/wk/{ => ui}/tmux.py | 0 7 files changed, 9 insertions(+), 7 deletions(-) rename scripts/wk/{ => ui}/tmux.py (100%) diff --git a/scripts/wk/__init__.py b/scripts/wk/__init__.py index 5bf5d701..afb3e8d5 100644 --- a/scripts/wk/__init__.py +++ b/scripts/wk/__init__.py @@ -17,7 +17,6 @@ from . import os from . import repairs from . import setup from . import std -from . import tmux from . import ui diff --git a/scripts/wk/clone/ddrescue.py b/scripts/wk/clone/ddrescue.py index d0255360..1c5b7cdc 100644 --- a/scripts/wk/clone/ddrescue.py +++ b/scripts/wk/clone/ddrescue.py @@ -20,7 +20,7 @@ from docopt import docopt import psutil import pytz -from wk import cfg, debug, exe, io, log, net, std, tmux +from wk import cfg, debug, exe, io, log, net, std from wk.cfg.ddrescue import ( DDRESCUE_MAP_TEMPLATE, DDRESCUE_SETTINGS, @@ -34,6 +34,7 @@ from wk.hw.smart import ( update_smart_details, ) from wk.ui import cli as ui # TODO: This is lazy +from wk.ui import tmux # STATIC VARIABLES diff --git a/scripts/wk/hw/cpu.py b/scripts/wk/hw/cpu.py index 826e0fb9..cbb04ba4 100644 --- a/scripts/wk/hw/cpu.py +++ b/scripts/wk/hw/cpu.py @@ -11,12 +11,12 @@ from wk import exe from wk.cfg.hw import CPU_FAILURE_TEMP from wk.os.mac import set_fans as macos_set_fans from wk.std import PLATFORM -from wk.tmux import respawn_pane as tmux_respawn_pane from wk.ui.cli import ( # TODO: This is lazy color_string, print_error, print_warning, ) +from wk.ui.tmux import respawn_pane as tmux_respawn_pane # STATIC VARIABLES diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py index 6c7ead59..1a3008f8 100644 --- a/scripts/wk/hw/diags.py +++ b/scripts/wk/hw/diags.py @@ -10,7 +10,7 @@ import time from docopt import docopt -from wk import cfg, debug, exe, log, std, tmux +from wk import cfg, debug, exe, log, std from wk.cfg.hw import STATUS_COLORS from wk.hw import benchmark as hw_benchmark from wk.hw import cpu as hw_cpu @@ -26,6 +26,7 @@ from wk.hw.screensavers import screensaver from wk.hw.test import Test, TestGroup from wk.ui import cli as ui # TODO: This is lazy +from wk.ui import tmux # STATIC VARIABLES diff --git a/scripts/wk/hw/screensavers.py b/scripts/wk/hw/screensavers.py index 4417777b..71e67fcd 100644 --- a/scripts/wk/hw/screensavers.py +++ b/scripts/wk/hw/screensavers.py @@ -6,7 +6,7 @@ import logging from subprocess import PIPE from wk.exe import run_program -from wk.tmux import zoom_pane as tmux_zoom_pane +from wk.ui import tmux # STATIC VARIABLES @@ -31,9 +31,9 @@ def screensaver(name) -> None: ] # Switch pane to fullscreen and start screensaver - tmux_zoom_pane() + tmux.zoom_pane() run_program(cmd, check=False, pipe=False, stderr=PIPE) - tmux_zoom_pane() + tmux.zoom_pane() if __name__ == '__main__': diff --git a/scripts/wk/ui/__init__.py b/scripts/wk/ui/__init__.py index eb6fa151..209a7699 100644 --- a/scripts/wk/ui/__init__.py +++ b/scripts/wk/ui/__init__.py @@ -1,3 +1,4 @@ """WizardKit: ui module init""" from . import cli +from . import tmux diff --git a/scripts/wk/tmux.py b/scripts/wk/ui/tmux.py similarity index 100% rename from scripts/wk/tmux.py rename to scripts/wk/ui/tmux.py From 03a143488ce659a6579a2dd5b262e2d083df5ef8 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sun, 2 Apr 2023 20:12:18 -0700 Subject: [PATCH 03/16] Adjust ui imports and calls --- scripts/wk/clone/ddrescue.py | 14 +++--- scripts/wk/graph.py | 14 +++--- scripts/wk/hw/benchmark.py | 11 ++--- scripts/wk/hw/cpu.py | 18 +++----- scripts/wk/hw/diags.py | 2 +- scripts/wk/hw/disk.py | 12 ++--- scripts/wk/hw/keyboard.py | 4 +- scripts/wk/hw/network.py | 15 +++---- scripts/wk/hw/sensors.py | 8 ++-- scripts/wk/hw/smart.py | 18 ++++---- scripts/wk/hw/surface_scan.py | 12 ++--- scripts/wk/hw/system.py | 6 +-- scripts/wk/kit/build_win.py | 31 +++++-------- scripts/wk/kit/ufd.py | 2 +- scripts/wk/net.py | 4 +- scripts/wk/os/linux.py | 8 ++-- scripts/wk/os/win.py | 14 +++--- scripts/wk/repairs/win.py | 81 ++++++++++++++------------------- scripts/wk/setup/win.py | 84 ++++++++++++++--------------------- 19 files changed, 154 insertions(+), 204 deletions(-) diff --git a/scripts/wk/clone/ddrescue.py b/scripts/wk/clone/ddrescue.py index 1c5b7cdc..1499d296 100644 --- a/scripts/wk/clone/ddrescue.py +++ b/scripts/wk/clone/ddrescue.py @@ -33,7 +33,7 @@ from wk.hw.smart import ( smart_status_ok, update_smart_details, ) -from wk.ui import cli as ui # TODO: This is lazy +from wk.ui import cli as ui from wk.ui import tmux @@ -1446,7 +1446,9 @@ def build_settings_menu(silent=True): preset = 'Default' if not silent: # Ask which preset to use - print(f'Available ddrescue presets: {" / ".join(SETTING_PRESETS)}') + ui.print_standard( + f'Available ddrescue presets: {" / ".join(SETTING_PRESETS)}' + ) preset = ui.choice(SETTING_PRESETS, 'Please select a preset:') # Fix selection @@ -2161,10 +2163,10 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): warning_message = 'Error(s) encountered, see message above' state.update_top_panes() if warning_message: - print(' ') - print(' ') + ui.print_standard(' ') + ui.print_standard(' ') ui.print_error('DDRESCUE PROCESS HALTED') - print(' ') + ui.print_standard(' ') ui.print_warning(warning_message) # Needs attention? @@ -2377,7 +2379,7 @@ def select_disk_parts(prompt, disk): object_list = [disk.path] # Convert object_list to hw_disk.Disk() objects - print(' ') + ui.print_standard(' ') ui.print_info('Getting disk/partition details...') object_list = [hw_disk.Disk(path) for path in object_list] diff --git a/scripts/wk/graph.py b/scripts/wk/graph.py index fe717545..0de162d5 100644 --- a/scripts/wk/graph.py +++ b/scripts/wk/graph.py @@ -3,7 +3,7 @@ import logging -from wk.ui.cli import color_string # TODO: This is lazy +from wk.ui import cli as ui # STATIC VARIABLES @@ -52,27 +52,27 @@ def generate_horizontal_graph(rate_list, graph_width=40, oneline=False): rate_color = 'GREEN' # Build graph - full_block = color_string((GRAPH_HORIZONTAL[-1],), (rate_color,)) + full_block = ui.color_string((GRAPH_HORIZONTAL[-1],), (rate_color,)) if step >= 24: - graph[0] += color_string((GRAPH_HORIZONTAL[step-24],), (rate_color,)) + graph[0] += ui.color_string((GRAPH_HORIZONTAL[step-24],), (rate_color,)) graph[1] += full_block graph[2] += full_block graph[3] += full_block elif step >= 16: graph[0] += ' ' - graph[1] += color_string((GRAPH_HORIZONTAL[step-16],), (rate_color,)) + graph[1] += ui.color_string((GRAPH_HORIZONTAL[step-16],), (rate_color,)) graph[2] += full_block graph[3] += full_block elif step >= 8: graph[0] += ' ' graph[1] += ' ' - graph[2] += color_string((GRAPH_HORIZONTAL[step-8],), (rate_color,)) + graph[2] += ui.color_string((GRAPH_HORIZONTAL[step-8],), (rate_color,)) graph[3] += full_block else: graph[0] += ' ' graph[1] += ' ' graph[2] += ' ' - graph[3] += color_string((GRAPH_HORIZONTAL[step],), (rate_color,)) + graph[3] += ui.color_string((GRAPH_HORIZONTAL[step],), (rate_color,)) # Done if oneline: @@ -128,7 +128,7 @@ def vertical_graph_line(percent, rate, scale=32): color_rate = 'GREEN' # Build string - line = color_string( + line = ui.color_string( strings=( f'{percent:5.1f}%', f'{GRAPH_VERTICAL[step]:<4}', diff --git a/scripts/wk/hw/benchmark.py b/scripts/wk/hw/benchmark.py index b336da5c..58d3ee83 100644 --- a/scripts/wk/hw/benchmark.py +++ b/scripts/wk/hw/benchmark.py @@ -22,10 +22,7 @@ from wk.cfg.hw import ( ) from wk.exe import run_program from wk.std import PLATFORM -from wk.ui.cli import ( # TODO: This is lazy - strip_colors, - color_string, - ) +from wk.ui import cli as ui # STATIC VARIABLES @@ -116,7 +113,7 @@ def check_io_results(test_obj, rate_list, graph_width) -> None: # Add horizontal graph to report for line in graph.generate_horizontal_graph(rate_list, graph_width): - if not strip_colors(line).strip(): + if not ui.strip_colors(line).strip(): # Skip empty lines continue test_obj.report.append(line) @@ -154,7 +151,7 @@ def run_io_test(test_obj, log_path, test_mode=False) -> None: LOG.info('Using %s for better performance', dev_path) offset = 0 read_rates = [] - test_obj.report.append(color_string('I/O Benchmark', 'BLUE')) + test_obj.report.append(ui.color_string('I/O Benchmark', 'BLUE')) # Get dd values or bail try: @@ -162,7 +159,7 @@ def run_io_test(test_obj, log_path, test_mode=False) -> None: except DeviceTooSmallError: test_obj.set_status('N/A') test_obj.report.append( - color_string('Disk too small to test', 'YELLOW'), + ui.color_string('Disk too small to test', 'YELLOW'), ) return diff --git a/scripts/wk/hw/cpu.py b/scripts/wk/hw/cpu.py index cbb04ba4..9a741ecd 100644 --- a/scripts/wk/hw/cpu.py +++ b/scripts/wk/hw/cpu.py @@ -11,12 +11,8 @@ from wk import exe from wk.cfg.hw import CPU_FAILURE_TEMP from wk.os.mac import set_fans as macos_set_fans from wk.std import PLATFORM -from wk.ui.cli import ( # TODO: This is lazy - color_string, - print_error, - print_warning, - ) -from wk.ui.tmux import respawn_pane as tmux_respawn_pane +from wk.ui import cli as ui +from wk.ui import tmux # STATIC VARIABLES @@ -97,9 +93,9 @@ def check_mprime_results(test_obj, working_dir) -> None: for line in passing_lines: test_obj.report.append(f' {line}') for line in warning_lines: - test_obj.report.append(color_string(f' {line}', 'YELLOW')) + test_obj.report.append(ui.color_string(f' {line}', 'YELLOW')) if not (passing_lines or warning_lines): - test_obj.report.append(color_string(' Unknown result', 'YELLOW')) + test_obj.report.append(ui.color_string(' Unknown result', 'YELLOW')) def start_mprime(working_dir, log_path) -> subprocess.Popen: @@ -147,7 +143,7 @@ def start_sysbench(sensors, sensors_out, log_path, pane) -> SysbenchType: ) # Update bottom pane - tmux_respawn_pane(pane, watch_file=log_path, watch_cmd='tail') + tmux.respawn_pane(pane, watch_file=log_path, watch_cmd='tail') # Start sysbench filehandle_sysbench = open( @@ -174,9 +170,9 @@ def set_apple_fan_speed(speed) -> None: except (RuntimeError, ValueError, subprocess.CalledProcessError) as err: LOG.error('Failed to set fans to %s', speed) LOG.error('Error: %s', err) - print_error(f'Failed to set fans to {speed}') + ui.print_error(f'Failed to set fans to {speed}') for line in str(err).splitlines(): - print_warning(f' {line.strip()}') + ui.print_warning(f' {line.strip()}') elif PLATFORM == 'Linux': cmd = ['apple-fans', speed] exe.run_program(cmd, check=False) diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py index 1a3008f8..72c594e4 100644 --- a/scripts/wk/hw/diags.py +++ b/scripts/wk/hw/diags.py @@ -25,7 +25,7 @@ from wk.hw.network import network_test from wk.hw.screensavers import screensaver from wk.hw.test import Test, TestGroup -from wk.ui import cli as ui # TODO: This is lazy +from wk.ui import cli as ui from wk.ui import tmux diff --git a/scripts/wk/hw/disk.py b/scripts/wk/hw/disk.py index 13bd821a..26a3d0c5 100644 --- a/scripts/wk/hw/disk.py +++ b/scripts/wk/hw/disk.py @@ -20,7 +20,7 @@ from wk.hw.smart import ( get_known_disk_attributes, ) from wk.std import PLATFORM -from wk.ui.cli import color_string, strip_colors # TODO: This is lazy +from wk.ui import cli as ui # STATIC VARIABLES @@ -74,7 +74,7 @@ class Disk: def add_note(self, note, color=None) -> None: """Add note that will be included in the disk report.""" if color: - note = color_string(note, color) + note = ui.color_string(note, color) if note not in self.notes: self.notes.append(note) self.notes.sort() @@ -83,7 +83,7 @@ class Disk: """Check if note is already present.""" present = False for note in self.notes: - if note_str == strip_colors(note): + if note_str == ui.strip_colors(note): present = True return present @@ -99,18 +99,18 @@ class Disk: """Generate Disk report, returns list.""" report = [] if header: - report.append(color_string(f'Device ({self.path.name})', 'BLUE')) + report.append(ui.color_string(f'Device ({self.path.name})', 'BLUE')) report.append(f' {self.description}') # Attributes if self.attributes: if header: - report.append(color_string('Attributes', 'BLUE')) + report.append(ui.color_string('Attributes', 'BLUE')) report.extend(generate_attribute_report(self)) # Notes if self.notes: - report.append(color_string('Notes', 'BLUE')) + report.append(ui.color_string('Notes', 'BLUE')) for note in self.notes: report.append(f' {note}') diff --git a/scripts/wk/hw/keyboard.py b/scripts/wk/hw/keyboard.py index 686fd52b..83541c45 100644 --- a/scripts/wk/hw/keyboard.py +++ b/scripts/wk/hw/keyboard.py @@ -5,7 +5,7 @@ import logging from wk.exe import run_program from wk.std import PLATFORM -from wk.ui.cli import print_warning # TODO: This is lazy +from wk.ui import cli as ui # STATIC VARIABLES @@ -18,7 +18,7 @@ def keyboard_test() -> None: if PLATFORM == 'Linux': run_xev() else: - print_warning(f'Not supported under this OS: {PLATFORM}') + ui.print_warning(f'Not supported under this OS: {PLATFORM}') def run_xev() -> None: diff --git a/scripts/wk/hw/network.py b/scripts/wk/hw/network.py index 29a7f8ba..55d02eb2 100644 --- a/scripts/wk/hw/network.py +++ b/scripts/wk/hw/network.py @@ -9,12 +9,7 @@ from wk.net import ( show_valid_addresses, speedtest, ) -from wk.ui.cli import ( - # TODO: This is lazy - TryAndPrint, - pause, - print_warning, - ) +from wk.ui import cli as ui # STATIC VARIABLES @@ -25,7 +20,7 @@ LOG = logging.getLogger(__name__) def network_test() -> None: """Run network tests.""" LOG.info('Network Test') - try_and_print = TryAndPrint() + try_and_print = ui.TryAndPrint() result = try_and_print.run( message='Network connection...', function=connected_to_private_network, @@ -35,8 +30,8 @@ def network_test() -> None: # Bail if not connected if result['Failed']: - print_warning('Please connect to a network and try again') - pause('Press Enter to return to main menu...') + ui.print_warning('Please connect to a network and try again') + ui.pause('Press Enter to return to main menu...') return # Show IP address(es) @@ -52,7 +47,7 @@ def network_test() -> None: try_and_print.run('Speedtest...', speedtest) # Done - pause('Press Enter to return to main menu...') + ui.pause('Press Enter to return to main menu...') if __name__ == '__main__': diff --git a/scripts/wk/hw/sensors.py b/scripts/wk/hw/sensors.py index 3ce1ded1..858b6553 100644 --- a/scripts/wk/hw/sensors.py +++ b/scripts/wk/hw/sensors.py @@ -13,7 +13,7 @@ from wk.cfg.hw import CPU_CRITICAL_TEMP, SMC_IDS, TEMP_COLORS from wk.exe import run_program, start_thread from wk.io import non_clobber_path from wk.std import PLATFORM, sleep -from wk.ui.cli import color_string # TODO: This is lazy +from wk.ui import cli as ui # STATIC VARIABLES @@ -110,7 +110,7 @@ class Sensors(): # Handle empty reports if not report: report = [ - color_string('WARNING: No sensors found', 'YELLOW'), + ui.color_string('WARNING: No sensors found', 'YELLOW'), '', 'Please monitor temps manually', ] @@ -426,7 +426,7 @@ def get_temp_str(temp, colored=True) -> str: temp = float(temp) except (TypeError, ValueError): # Invalid temp? - return color_string(temp, 'PURPLE') + return ui.color_string(temp, 'PURPLE') # Determine color if colored: @@ -436,7 +436,7 @@ def get_temp_str(temp, colored=True) -> str: break # Done - return color_string(f'{"-" if temp < 0 else ""}{temp:2.0f}°C', temp_color) + return ui.color_string(f'{"-" if temp < 0 else ""}{temp:2.0f}°C', temp_color) diff --git a/scripts/wk/hw/smart.py b/scripts/wk/hw/smart.py index e56db1ea..c777f400 100644 --- a/scripts/wk/hw/smart.py +++ b/scripts/wk/hw/smart.py @@ -19,7 +19,7 @@ from wk.cfg.hw import ( ) from wk.exe import get_json_from_command, run_program from wk.std import bytes_to_string, sleep -from wk.ui.cli import color_string # TODO: This is lazy +from wk.ui import cli as ui # STATIC VARIABLES @@ -41,26 +41,26 @@ def build_self_test_report(test_obj, aborted=False) -> None: For instance if the test was aborted the report should include the last known progress instead of just "was aborted by host." """ - report = [color_string('Self-Test', 'BLUE')] + report = [ui.color_string('Self-Test', 'BLUE')] test_details = get_smart_self_test_details(test_obj.dev) test_result = test_details.get('status', {}).get('string', 'Unknown') # Build report if test_obj.disabled or test_obj.status == 'Denied': - report.append(color_string(f' {test_obj.status}', 'RED')) + report.append(ui.color_string(f' {test_obj.status}', 'RED')) elif test_obj.status == 'N/A' or not test_obj.dev.attributes: - report.append(color_string(f' {test_obj.status}', 'YELLOW')) + report.append(ui.color_string(f' {test_obj.status}', 'YELLOW')) elif test_obj.status == 'TestInProgress': - report.append(color_string(' Failed to stop previous test', 'RED')) + report.append(ui.color_string(' Failed to stop previous test', 'RED')) test_obj.set_status('Failed') else: # Other cases include self-test result string report.append(f' {test_result.capitalize()}') if aborted and not (test_obj.passed or test_obj.failed): - report.append(color_string(' Aborted', 'YELLOW')) + report.append(ui.color_string(' Aborted', 'YELLOW')) test_obj.set_status('Aborted') elif test_obj.status == 'TimedOut': - report.append(color_string(' TimedOut', 'YELLOW')) + report.append(ui.color_string(' TimedOut', 'YELLOW')) # Done test_obj.report.extend(report) @@ -137,7 +137,7 @@ def generate_attribute_report(dev, only_failed=False) -> list[str]: continue # Build colored string and append to report - line = color_string( + line = ui.color_string( [label, get_attribute_value_string(dev, attr), note], [None, value_color, 'YELLOW'], ) @@ -299,7 +299,7 @@ def run_smart_self_test(test_obj, log_path) -> bool: finished = False test_details = get_smart_self_test_details(test_obj.dev) size_str = bytes_to_string(test_obj.dev.size, use_binary=False) - header_str = color_string( + header_str = ui.color_string( ['[', test_obj.dev.path.name, ' ', size_str, ']'], [None, 'BLUE', None, 'CYAN', None], sep='', diff --git a/scripts/wk/hw/surface_scan.py b/scripts/wk/hw/surface_scan.py index b2dc99c6..9a542d30 100644 --- a/scripts/wk/hw/surface_scan.py +++ b/scripts/wk/hw/surface_scan.py @@ -15,7 +15,7 @@ from wk.cfg.hw import ( ) from wk.exe import run_program from wk.std import PLATFORM, bytes_to_string -from wk.ui.cli import color_string, strip_colors # TODO: This is lazy +from wk.ui import cli as ui # STATIC VARIABLES @@ -27,7 +27,7 @@ def check_surface_scan_results(test_obj, log_path) -> None: """Check results and set test status.""" with open(log_path, 'r', encoding='utf-8') as _f: for line in _f.readlines(): - line = strip_colors(line.strip()) + line = ui.strip_colors(line.strip()) if not line or BADBLOCKS_SKIP_REGEX.match(line): # Skip continue @@ -44,10 +44,10 @@ def check_surface_scan_results(test_obj, log_path) -> None: test_obj.set_status('Passed') else: test_obj.failed = True - test_obj.report.append(f' {color_string(line, "YELLOW")}') + test_obj.report.append(f' {ui.color_string(line, "YELLOW")}') test_obj.set_status('Failed') else: - test_obj.report.append(f' {color_string(line, "YELLOW")}') + test_obj.report.append(f' {ui.color_string(line, "YELLOW")}') if not (test_obj.passed or test_obj.failed): test_obj.set_status('Unknown') @@ -61,7 +61,7 @@ def run_scan(test_obj, log_path, test_mode=False) -> None: # Use "RAW" disks under macOS dev_path = dev_path.with_name(f'r{dev_path.name}') LOG.info('Using %s for better performance', dev_path) - test_obj.report.append(color_string('badblocks', 'BLUE')) + test_obj.report.append(ui.color_string('badblocks', 'BLUE')) test_obj.set_status('Working') # Increase block size if necessary @@ -80,7 +80,7 @@ def run_scan(test_obj, log_path, test_mode=False) -> None: with open(log_path, 'a', encoding='utf-8') as _f: size_str = bytes_to_string(dev.size, use_binary=False) _f.write( - color_string( + ui.color_string( ['[', dev.path.name, ' ', size_str, ']\n'], [None, 'BLUE', None, 'CYAN', None], sep='', diff --git a/scripts/wk/hw/system.py b/scripts/wk/hw/system.py index 6634b326..b82f9433 100644 --- a/scripts/wk/hw/system.py +++ b/scripts/wk/hw/system.py @@ -13,7 +13,7 @@ from wk.cfg.python import DATACLASS_DECORATOR_KWARGS from wk.exe import get_json_from_command, run_program from wk.hw.test import Test from wk.std import PLATFORM, bytes_to_string, string_to_bytes -from wk.ui.cli import color_string # TODO: This is lazy +from wk.ui import cli as ui # STATIC VARIABLES @@ -37,11 +37,11 @@ class System: def generate_report(self) -> list[str]: """Generate CPU & RAM report, returns list.""" report = [] - report.append(color_string('Device', 'BLUE')) + report.append(ui.color_string('Device', 'BLUE')) report.append(f' {self.cpu_description}') # Include RAM details - report.append(color_string('RAM', 'BLUE')) + report.append(ui.color_string('RAM', 'BLUE')) report.append(f' {self.ram_total} ({", ".join(self.ram_dimms)})') # Tests diff --git a/scripts/wk/kit/build_win.py b/scripts/wk/kit/build_win.py index bb363430..e87d5101 100644 --- a/scripts/wk/kit/build_win.py +++ b/scripts/wk/kit/build_win.py @@ -22,16 +22,7 @@ from wk.kit.tools import ( ) from wk.log import update_log_path from wk.std import GenericError -from wk.ui.cli import ( - # TODO: This is lazy - TryAndPrint, - clear_screen, - pause, - print_info, - print_success, - set_title, - sleep, - ) +from wk.ui import cli as ui # STATIC VARIABLES @@ -246,7 +237,7 @@ def download_libreoffice(): for arch in 32, 64: out_path = INSTALLERS_DIR.joinpath(f'LibreOffice{arch}.msi') download_file(out_path, SOURCES[f'LibreOffice{arch}']) - sleep(1) + ui.sleep(1) def download_neutron(): @@ -316,7 +307,7 @@ def download_snappy_driver_installer_origin(): cmd.append('-new_console:n') cmd.append('-new_console:s33V') popen_program(cmd, cwd=aria2c.parent) - sleep(1) + ui.sleep(1) wait_for_procs('aria2c.exe') else: run_program(cmd) @@ -460,13 +451,13 @@ def build_kit(): """Build Kit.""" update_log_path(dest_name='Build Tool', timestamp=True) title = f'{KIT_NAME_FULL}: Build Tool' - clear_screen() - set_title(title) - print_info(title) + ui.clear_screen() + ui.set_title(title) + ui.print_info(title) print('') # Set up TryAndPrint - try_print = TryAndPrint() + try_print = ui.TryAndPrint() try_print.width = WIDTH try_print.verbose = True for error in ('CalledProcessError', 'FileNotFoundError'): @@ -497,15 +488,15 @@ def build_kit(): # Pause print('', flush=True) - pause('Please review and press Enter to continue...') + ui.pause('Please review and press Enter to continue...') # Compress .cbin try_print.run('Compress cbin...', compress_cbin_dirs) # Generate launcher scripts - print_success('Generating launchers') + ui.print_success('Generating launchers') for section, launchers in sorted(LAUNCHERS.items()): - print_info(f' {section if section else "(Root)"}') + ui.print_info(f' {section if section else "(Root)"}') for name, options in sorted(launchers.items()): try_print.run( f' {name}...', generate_launcher, @@ -515,7 +506,7 @@ def build_kit(): # Done print('') print('Done.') - pause('Press Enter to exit...') + ui.pause('Press Enter to exit...') if __name__ == '__main__': diff --git a/scripts/wk/kit/ufd.py b/scripts/wk/kit/ufd.py index f0cfa767..64697a41 100644 --- a/scripts/wk/kit/ufd.py +++ b/scripts/wk/kit/ufd.py @@ -23,7 +23,7 @@ from wk.cfg.ufd import ( from wk.exe import get_json_from_command, run_program from wk.os import linux -from wk.ui import cli as ui # TODO: This is lazy +from wk.ui import cli as ui # STATIC VARIABLES diff --git a/scripts/wk/net.py b/scripts/wk/net.py index 1c4c3383..5e478e54 100644 --- a/scripts/wk/net.py +++ b/scripts/wk/net.py @@ -11,7 +11,7 @@ from wk.exe import get_json_from_command, run_program from wk.std import PLATFORM, GenericError from wk.cfg.net import BACKUP_SERVERS -from wk.ui.cli import show_data # TODO: This is lazy +from wk.ui import cli as ui # REGEX @@ -200,7 +200,7 @@ def show_valid_addresses(): for family in families: if REGEX_VALID_IP.search(family.address): # Valid IP found - show_data(message=dev, data=family.address) + ui.show_data(message=dev, data=family.address) def speedtest(): diff --git a/scripts/wk/os/linux.py b/scripts/wk/os/linux.py index 1c134f78..991dd454 100644 --- a/scripts/wk/os/linux.py +++ b/scripts/wk/os/linux.py @@ -11,7 +11,7 @@ from wk.cfg.hw import VOLUME_FAILURE_THRESHOLD, VOLUME_WARNING_THRESHOLD from wk.exe import get_json_from_command, popen_program, run_program from wk.log import format_log_path from wk.std import bytes_to_string -from wk.ui.cli import color_string # TODO: This is lazy +from wk.ui import cli as ui # STATIC VARIABLES @@ -83,20 +83,20 @@ def build_volume_report(device_path=None) -> list: vol['mountpoint'] = f'Mounted on {vol["mountpoint"]}' # Name and size - line = color_string( + line = ui.color_string( [f'{vol["name"]:<20}', f'{vol["size"]:>9}'], [None, 'CYAN'], ) # Mountpoint and type - line = color_string( + line = ui.color_string( [line, f'{vol["mountpoint"]:<{m_width}}', f'{vol["fstype"]:<11}'], [None, None, 'BLUE'], ) # Used and free if any([vol['fsused'], vol['fsavail']]): - line = color_string( + line = ui.color_string( [line, f'({vol["fsused"]:>9} used, {vol["fsavail"]:>9} free)'], [None, size_color], ) diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py index 8e7a3dc1..852f6065 100644 --- a/scripts/wk/os/win.py +++ b/scripts/wk/os/win.py @@ -32,7 +32,7 @@ from wk.std import ( bytes_to_string, sleep, ) -from wk.ui.cli import color_string, input_text # TODO: This is lazy +from wk.ui import cli as ui # STATIC VARIABLES @@ -182,7 +182,7 @@ def check_4k_alignment(show_alert=False): continue if int(match.group('offset')) % 4096 != 0: report.append( - color_string( + ui.color_string( f'{match.group("description")}' f' ({bytes_to_string(match.group("size"), decimals=1)})' , @@ -198,7 +198,7 @@ def check_4k_alignment(show_alert=False): if report: report.insert( 0, - color_string('One or more partitions not 4K aligned', 'YELLOW'), + ui.color_string('One or more partitions not 4K aligned', 'YELLOW'), ) return report @@ -211,7 +211,7 @@ def export_bitlocker_info(): ] # Get filename - file_name = input_text(prompt='Enter filename', allow_empty_response=False) + file_name = ui.input_text(prompt='Enter filename', allow_empty_response=False) file_path = pathlib.Path(f'../../Bitlocker_{file_name}.txt').resolve() # Save info @@ -250,13 +250,13 @@ def get_installed_antivirus(): state = proc.stdout.split('=')[1] state = hex(int(state)) if str(state)[3:5] not in ['10', '11']: - report.append(color_string(f'[Disabled] {product}', 'YELLOW')) + report.append(ui.color_string(f'[Disabled] {product}', 'YELLOW')) else: report.append(product) # Final check if not report: - report.append(color_string('No products detected', 'RED')) + report.append(ui.color_string('No products detected', 'RED')) # Done return report @@ -363,7 +363,7 @@ def get_volume_usage(use_colors=False): f' ({bytes_to_string(free, 2):>10} / {bytes_to_string(total, 2):>10})' ) if use_colors: - display_str = color_string(display_str, color) + display_str = ui.color_string(display_str, color) report.append(f'{disk.device} {display_str}') # Done diff --git a/scripts/wk/repairs/win.py b/scripts/wk/repairs/win.py index f468be33..d94f6e3d 100644 --- a/scripts/wk/repairs/win.py +++ b/scripts/wk/repairs/win.py @@ -60,22 +60,7 @@ from wk.std import ( GenericWarning, sleep, ) -from wk.ui.cli import ( - # TODO: This is lazy - Menu, - TryAndPrint, - abort, - ask, - clear_screen, - color_string, - pause, - print_info, - print_standard, - print_warning, - set_title, - show_data, - strip_colors, - ) +from wk.ui import cli as ui # STATIC VARIABLES @@ -90,7 +75,7 @@ GPUPDATE_SUCCESS_STRINGS = ( 'User Policy update has completed successfully.', ) IN_CONEMU = 'ConEmuPID' in os.environ -MENU_PRESETS = Menu() +MENU_PRESETS = ui.Menu() PROGRAMDATA = os.environ.get('{ALLUSERSPROFILE}', r'C:\ProgramData') PROGRAMFILES_32 = os.environ.get( 'PROGRAMFILES(X86)', os.environ.get( @@ -108,7 +93,7 @@ WHITELIST = '\n'.join(( fr'{PROGRAMFILES_32}\TeamViewer\tv_x64.exe', sys.executable, )) -TRY_PRINT = TryAndPrint() +TRY_PRINT = ui.TryAndPrint() TRY_PRINT.width = WIDTH TRY_PRINT.verbose = True for error in ('CalledProcessError', 'FileNotFoundError'): @@ -119,7 +104,7 @@ for error in ('CalledProcessError', 'FileNotFoundError'): def build_menus(base_menus, title, presets): """Build menus, returns dict.""" menus = {} - menus['Main'] = Menu(title=f'{title}\n{color_string("Main Menu", "GREEN")}') + menus['Main'] = ui.Menu(title=f'{title}\n{ui.color_string("Main Menu", "GREEN")}') # Main Menu for entry in base_menus['Actions']: @@ -128,7 +113,7 @@ def build_menus(base_menus, title, presets): menus['Main'].add_option(group, {'Selected': True}) # Options - menus['Options'] = Menu(title=f'{title}\n{color_string("Options", "GREEN")}') + menus['Options'] = ui.Menu(title=f'{title}\n{ui.color_string("Options", "GREEN")}') for entry in base_menus['Options']: menus['Options'].add_option(entry.name, entry.details) menus['Options'].add_action('All') @@ -138,7 +123,7 @@ def build_menus(base_menus, title, presets): # Run groups for group, entries in base_menus['Groups'].items(): - menus[group] = Menu(title=f'{title}\n{color_string(group, "GREEN")}') + menus[group] = ui.Menu(title=f'{title}\n{ui.color_string(group, "GREEN")}') menus[group].disabled_str = 'Locked' for entry in entries: menus[group].add_option(entry.name, entry.details) @@ -170,7 +155,7 @@ def build_menus(base_menus, title, presets): ) # Update presets Menu - MENU_PRESETS.title = f'{title}\n{color_string("Load Preset", "GREEN")}' + MENU_PRESETS.title = f'{title}\n{ui.color_string("Load Preset", "GREEN")}' MENU_PRESETS.add_option('Default') for name in presets: MENU_PRESETS.add_option(name) @@ -274,7 +259,7 @@ def init(menus, presets): # Resume session load_settings(menus) - print_info('Resuming session, press CTRL+c to cancel') + ui.print_info('Resuming session, press CTRL+c to cancel') for _x in range(AUTO_REPAIR_DELAY_IN_SECONDS, 0, -1): print(f' {_x} second{"" if _x==1 else "s"} remaining... \r', end='') sleep(1) @@ -319,7 +304,7 @@ def init_session(options): 'The timezone is currently set to ' f'{zone}, switch it to {WINDOWS_TIME_ZONE}?' ) - if zone != WINDOWS_TIME_ZONE and ask(msg): + if zone != WINDOWS_TIME_ZONE and ui.ask(msg): set_timezone(WINDOWS_TIME_ZONE) # One-time tasks @@ -395,20 +380,20 @@ def load_settings(menus): if group == 'Main': continue for name in menu.options: - menu.options[name].update(get_entry_settings(group, strip_colors(name))) + menu.options[name].update(get_entry_settings(group, ui.strip_colors(name))) def run_auto_repairs(base_menus, presets): """Run Auto Repairs.""" set_log_path() title = f'{KIT_NAME_FULL}: Auto Repairs' - clear_screen() - set_title(title) - print_info(title) + ui.clear_screen() + ui.set_title(title) + ui.print_info(title) print('') # Generate menus - print_standard('Initializing...') + ui.print_standard('Initializing...') menus = build_menus(base_menus, title, presets) # Init @@ -426,21 +411,21 @@ def run_auto_repairs(base_menus, presets): try: show_main_menu(base_menus, menus, presets, title) except SystemExit: - if ask('End session?'): + if ui.ask('End session?'): end_session() raise # Start or resume repairs - clear_screen() - print_standard(title) + ui.clear_screen() + ui.print_standard(title) print('') save_selection_settings(menus) - print_info('Initializing...') + ui.print_info('Initializing...') init_run(menus['Options'].options) save_selection_settings(menus) if not session_started: init_session(menus['Options'].options) - print_info('Running repairs') + ui.print_info('Running repairs') # Run repairs for group, menu in menus.items(): @@ -449,19 +434,19 @@ def run_auto_repairs(base_menus, presets): try: run_group(group, menu) except KeyboardInterrupt: - abort() + ui.abort() # Done end_session() - print_info('Done') - pause('Press Enter to exit...') + ui.print_info('Done') + ui.pause('Press Enter to exit...') def run_group(group, menu): """Run entries in group if appropriate.""" - print_info(f' {group}') + ui.print_info(f' {group}') for name, details in menu.options.items(): - name_str = strip_colors(name) + name_str = ui.strip_colors(name) skipped = details.get('Skipped', False) done = details.get('Done', False) disabled = details.get('Disabled', False) @@ -475,7 +460,7 @@ def run_group(group, menu): # Previously skipped if skipped: - show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH) + ui.show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH) continue # Previously ran @@ -485,7 +470,7 @@ def run_group(group, menu): color = 'YELLOW' elif details.get('Failed', False): color = 'RED' - show_data( + ui.show_data( f'{name_str}...', details.get('Message', 'Unknown'), color, width=WIDTH, ) @@ -493,7 +478,7 @@ def run_group(group, menu): # Not selected if not selected: - show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH) + ui.show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH) save_settings(group, name, skipped=True) continue @@ -516,7 +501,7 @@ def save_selection_settings(menus): def save_settings(group, name, result=None, **kwargs): """Save entry settings in the registry.""" - key_path = fr'{AUTO_REPAIR_KEY}\{group}\{strip_colors(name)}' + key_path = fr'{AUTO_REPAIR_KEY}\{group}\{ui.strip_colors(name)}' # Get values from TryAndPrint result if result: @@ -530,7 +515,7 @@ def save_settings(group, name, result=None, **kwargs): # Write values to registry for value_name, data in kwargs.items(): - value_name = strip_colors(value_name) + value_name = ui.strip_colors(value_name) if isinstance(data, bool): data = 1 if data else 0 if isinstance(data, int): @@ -894,7 +879,7 @@ def backup_all_browser_profiles(use_try_print=False): users = get_path_obj(f'{SYSTEMDRIVE}/Users') for userprofile in users.iterdir(): if use_try_print: - print_info(f'{" "*6}{userprofile.name}') + ui.print_info(f'{" "*6}{userprofile.name}') backup_browser_profiles(userprofile, use_try_print) @@ -908,7 +893,7 @@ def backup_browser_chromium(backup_path, browser, search_path, use_try_print): if output_path.exists(): # Assuming backup was already done if use_try_print: - show_data( + ui.show_data( f'{" "*8}{browser} ({item.name})...', 'Backup already exists.', color='YELLOW', width=WIDTH, ) @@ -936,7 +921,7 @@ def backup_browser_firefox(backup_path, search_path, use_try_print): if output_path.exists(): # Assuming backup was already done if use_try_print: - show_data( + ui.show_data( f'{" "*8}Firefox (All)...', 'Backup already exists.', color='YELLOW', width=WIDTH, ) @@ -1322,7 +1307,7 @@ def kill_explorer(): def reboot(timeout=10): """Reboot the system.""" atexit.unregister(start_explorer) - print_warning(f'Rebooting the system in {timeout} seconds...') + ui.print_warning(f'Rebooting the system in {timeout} seconds...') sleep(timeout) cmd = ['shutdown', '-r', '-t', '0'] run_program(cmd, check=False) diff --git a/scripts/wk/setup/win.py b/scripts/wk/setup/win.py index abf0e0db..22d8224f 100644 --- a/scripts/wk/setup/win.py +++ b/scripts/wk/setup/win.py @@ -62,23 +62,7 @@ from wk.std import ( GenericWarning, sleep, ) -from wk.ui.cli import ( - # TODO: This is lazy - Menu, - TryAndPrint, - abort, - ask, - clear_screen, - color_string, - pause, - print_error, - print_info, - print_standard, - print_warning, - set_title, - show_data, - strip_colors, - ) +from wk.ui import cli as ui # STATIC VARIABLES @@ -94,7 +78,7 @@ KNOWN_ENCODINGS = ( 'utf-32-le', ) IN_CONEMU = 'ConEmuPID' in os.environ -MENU_PRESETS = Menu() +MENU_PRESETS = ui.Menu() PROGRAMFILES_32 = os.environ.get( 'PROGRAMFILES(X86)', os.environ.get( 'PROGRAMFILES', r'C:\Program Files (x86)', @@ -106,7 +90,7 @@ PROGRAMFILES_64 = os.environ.get( ), ) SYSTEMDRIVE = os.environ.get('SYSTEMDRIVE', 'C:') -TRY_PRINT = TryAndPrint() +TRY_PRINT = ui.TryAndPrint() TRY_PRINT.width = WIDTH TRY_PRINT.verbose = True for error in ('CalledProcessError', 'FileNotFoundError'): @@ -117,7 +101,7 @@ for error in ('CalledProcessError', 'FileNotFoundError'): def build_menus(base_menus, title, presets): """Build menus, returns dict.""" menus = {} - menus['Main'] = Menu(title=f'{title}\n{color_string("Main Menu", "GREEN")}') + menus['Main'] = ui.Menu(title=f'{title}\n{ui.color_string("Main Menu", "GREEN")}') # Main Menu for entry in base_menus['Actions']: @@ -127,7 +111,7 @@ def build_menus(base_menus, title, presets): # Run groups for group, entries in base_menus['Groups'].items(): - menus[group] = Menu(title=f'{title}\n{color_string(group, "GREEN")}') + menus[group] = ui.Menu(title=f'{title}\n{ui.color_string(group, "GREEN")}') for entry in entries: menus[group].add_option(entry.name, entry.details) menus[group].add_action('All') @@ -156,7 +140,7 @@ def build_menus(base_menus, title, presets): ) # Update presets Menu - MENU_PRESETS.title = f'{title}\n{color_string("Load Preset", "GREEN")}' + MENU_PRESETS.title = f'{title}\n{ui.color_string("Load Preset", "GREEN")}' MENU_PRESETS.add_option('Default') for name in presets: MENU_PRESETS.add_option(name) @@ -173,26 +157,26 @@ def check_os_and_set_menu_title(title): """Check OS version and update title for menus, returns str.""" color = None os_name = get_os_name(check=False) - print_standard(f'Operating System: {os_name}') + ui.print_standard(f'Operating System: {os_name}') # Check support status and set color try: get_os_name() except GenericWarning: # Outdated version - print_warning('OS version is outdated, updating is recommended.') - if not ask('Continue anyway?'): - abort() + ui.print_warning('OS version is outdated, updating is recommended.') + if not ui.ask('Continue anyway?'): + ui.abort() color = 'YELLOW' except GenericError: # Unsupported version - print_error('OS version is unsupported, updating is recommended.') - if not ask('Continue anyway? (NOT RECOMMENDED)'): - abort() + ui.print_error('OS version is unsupported, updating is recommended.') + if not ui.ask('Continue anyway? (NOT RECOMMENDED)'): + ui.abort() color = 'RED' # Done - return f'{title} ({color_string(os_name, color)})' + return f'{title} ({ui.color_string(os_name, color)})' def load_preset(menus, presets, title, enable_menu_exit=True): @@ -218,10 +202,10 @@ def load_preset(menus, presets, title, enable_menu_exit=True): menu.options[name]['Selected'] = value # Ask selection question(s) - clear_screen() - print_standard(f'{title}') + ui.clear_screen() + ui.print_standard(f'{title}') print('') - if selection[0] == 'Default' and ask('Install LibreOffice?'): + if selection[0] == 'Default' and ui.ask('Install LibreOffice?'): menus['Install Software'].options['LibreOffice']['Selected'] = True # Re-enable Main Menu action if disabled @@ -238,11 +222,11 @@ def run_auto_setup(base_menus, presets): """Run Auto Setup.""" update_log_path(dest_name='Auto Setup', timestamp=True) title = f'{KIT_NAME_FULL}: Auto Setup' - clear_screen() - set_title(title) - print_info(title) + ui.clear_screen() + ui.set_title(title) + ui.print_info(title) print('') - print_standard('Initializing...') + ui.print_standard('Initializing...') # Check OS and update title for menus title = check_os_and_set_menu_title(title) @@ -257,10 +241,10 @@ def run_auto_setup(base_menus, presets): show_main_menu(base_menus, menus, presets, title) # Start setup - clear_screen() - print_standard(title) + ui.clear_screen() + ui.print_standard(title) print('') - print_info('Running setup') + ui.print_info('Running setup') # Run setup for group, menu in menus.items(): @@ -269,22 +253,22 @@ def run_auto_setup(base_menus, presets): try: run_group(group, menu) except KeyboardInterrupt: - abort() + ui.abort() # Done - print_info('Done') - pause('Press Enter to exit...') + ui.print_info('Done') + ui.pause('Press Enter to exit...') def run_group(group, menu): """Run entries in group if appropriate.""" - print_info(f' {group}') + ui.print_info(f' {group}') for name, details in menu.options.items(): - name_str = strip_colors(name) + name_str = ui.strip_colors(name) # Not selected if not details.get('Selected', False): - show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH) + ui.show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH) continue # Selected @@ -421,7 +405,7 @@ def auto_config_browsers(): 'Set default browser...', set_default_browser, msg_good='STARTED', ) print(prompt, end='', flush=True) - pause('') + ui.pause('') # Move cursor to beginning of the previous line and clear prompt print(f'\033[F\r{" "*len(prompt)}\r', end='', flush=True) @@ -807,8 +791,8 @@ def install_software_bundle(): warning = 'NOTE: Press CTRL+c to manually resume if it gets stuck...' # Start installations and wait for them to finish - print_standard(msg) - print_warning(warning, end='', flush=True) + ui.print_standard(msg) + ui.print_warning(warning, end='', flush=True) proc = popen_program([installer]) try: proc.wait() @@ -909,7 +893,7 @@ def get_storage_status(): """Get storage status for fixed disks, returns list.""" report = get_volume_usage(use_colors=True) for disk in get_raw_disks(): - report.append(color_string(f'Uninitialized Disk: {disk}', 'RED')) + report.append(ui.color_string(f'Uninitialized Disk: {disk}', 'RED')) # Done return report From 95d71594147aa0a971691d8b9fa9c38940cbfc6d Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sun, 2 Apr 2023 20:46:54 -0700 Subject: [PATCH 04/16] Move ANSI color escape sections to their own file --- scripts/auto_repairs.py | 2 +- scripts/mount-backup-shares | 2 +- scripts/unmount-backup-shares | 2 +- scripts/wk/__init__.py | 1 + scripts/wk/ansi.py | 67 ++++++++++++++++++++++++++++ scripts/wk/clone/ddrescue.py | 84 +++++++++++++++++------------------ scripts/wk/graph.py | 14 +++--- scripts/wk/hw/benchmark.py | 9 ++-- scripts/wk/hw/cpu.py | 6 +-- scripts/wk/hw/diags.py | 24 +++++----- scripts/wk/hw/disk.py | 12 ++--- scripts/wk/hw/sensors.py | 8 ++-- scripts/wk/hw/smart.py | 18 ++++---- scripts/wk/hw/surface_scan.py | 12 ++--- scripts/wk/hw/system.py | 6 +-- scripts/wk/os/linux.py | 8 ++-- scripts/wk/os/win.py | 11 ++--- scripts/wk/repairs/win.py | 17 +++---- scripts/wk/setup/win.py | 13 +++--- scripts/wk/ui/cli.py | 56 +---------------------- 20 files changed, 194 insertions(+), 178 deletions(-) create mode 100644 scripts/wk/ansi.py diff --git a/scripts/auto_repairs.py b/scripts/auto_repairs.py index 0615f115..b68f5e0b 100644 --- a/scripts/auto_repairs.py +++ b/scripts/auto_repairs.py @@ -5,7 +5,7 @@ import wk # Classes -REBOOT_STR = wk.ui.cli.color_string('Reboot', 'YELLOW') +REBOOT_STR = wk.ansi.color_string('Reboot', 'YELLOW') class MenuEntry(): """Simple class to allow cleaner code below.""" def __init__(self, name, function=None, selected=True, **kwargs): diff --git a/scripts/mount-backup-shares b/scripts/mount-backup-shares index bf528771..973b76e2 100755 --- a/scripts/mount-backup-shares +++ b/scripts/mount-backup-shares @@ -17,7 +17,7 @@ def main(): color = 'RED' elif 'Already' in line: color = 'YELLOW' - print(wk.ui.cli.color_string(line, color)) + print(wk.ansi.color_string(line, color)) if __name__ == '__main__': diff --git a/scripts/unmount-backup-shares b/scripts/unmount-backup-shares index 57e7e880..81188020 100755 --- a/scripts/unmount-backup-shares +++ b/scripts/unmount-backup-shares @@ -15,7 +15,7 @@ def main(): line = f' {line}' if 'Not mounted' in line: color = 'YELLOW' - print(wk.ui.cli.color_string(line, color)) + print(wk.ansi.color_string(line, color)) if __name__ == '__main__': diff --git a/scripts/wk/__init__.py b/scripts/wk/__init__.py index afb3e8d5..2a77ed94 100644 --- a/scripts/wk/__init__.py +++ b/scripts/wk/__init__.py @@ -3,6 +3,7 @@ from sys import stderr, version_info +from . import ansi from . import cfg from . import clone from . import debug diff --git a/scripts/wk/ansi.py b/scripts/wk/ansi.py new file mode 100644 index 00000000..2acf515d --- /dev/null +++ b/scripts/wk/ansi.py @@ -0,0 +1,67 @@ +"""WizardKit: ANSI control/escape functions""" +# vim: sts=2 sw=2 ts=2 + +import itertools +import logging +import pathlib + +# STATIC VARIABLES +LOG = logging.getLogger(__name__) +COLORS = { + 'CLEAR': '\033[0m', + 'RED': '\033[31m', + 'RED_BLINK': '\033[31;5m', + 'ORANGE': '\033[31;1m', + 'ORANGE_RED': '\033[1;31;41m', + 'GREEN': '\033[32m', + 'YELLOW': '\033[33m', + 'YELLOW_BLINK': '\033[33;5m', + 'BLUE': '\033[34m', + 'PURPLE': '\033[35m', + 'CYAN': '\033[36m', + } + + +# Functions +def color_string(strings, colors, sep=' '): + """Build colored string using ANSI escapes, returns str.""" + clear_code = COLORS['CLEAR'] + msg = [] + + # Convert to tuples if necessary + if isinstance(strings, (str, pathlib.Path)): + strings = (strings,) + if isinstance(colors, (str, pathlib.Path)): + colors = (colors,) + + # Convert to strings if necessary + try: + iter(strings) + except TypeError: + # Assuming single element passed, convert to string + strings = (str(strings),) + try: + iter(colors) + except TypeError: + # Assuming single element passed, convert to string + colors = (str(colors),) + + # Build new string with color escapes added + for string, color in itertools.zip_longest(strings, colors): + color_code = COLORS.get(color, clear_code) + msg.append(f'{color_code}{string}{clear_code}') + + # Done + return sep.join(msg) + + +def strip_colors(string): + """Strip known ANSI color escapes from string, returns str.""" + LOG.debug('string: %s', string) + for color in COLORS.values(): + string = string.replace(color, '') + return string + + +if __name__ == '__main__': + print("This file is not meant to be called directly.") diff --git a/scripts/wk/clone/ddrescue.py b/scripts/wk/clone/ddrescue.py index 1499d296..672c3e25 100644 --- a/scripts/wk/clone/ddrescue.py +++ b/scripts/wk/clone/ddrescue.py @@ -20,7 +20,7 @@ from docopt import docopt import psutil import pytz -from wk import cfg, debug, exe, io, log, net, std +from wk import ansi, cfg, debug, exe, io, log, net, std from wk.cfg.ddrescue import ( DDRESCUE_MAP_TEMPLATE, DDRESCUE_SETTINGS, @@ -91,8 +91,8 @@ REGEX_REMAINING_TIME = re.compile( LOG = logging.getLogger(__name__) MENU_ACTIONS = ( 'Start', - f'Change settings {ui.color_string("(experts only)", "YELLOW")}', - f'Detect drives {ui.color_string("(experts only)", "YELLOW")}', + f'Change settings {ansi.color_string("(experts only)", "YELLOW")}', + f'Detect drives {ansi.color_string("(experts only)", "YELLOW")}', 'Quit') MENU_TOGGLES = { 'Auto continue (if recovery % over threshold)': True, @@ -422,7 +422,7 @@ class State(): self.panes['Started'] = tmux.split_window( lines=cfg.ddrescue.TMUX_SIDE_WIDTH, target_id=self.panes['Source'], - text=ui.color_string( + text=ansi.color_string( ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], ['BLUE', None], sep='\n', @@ -568,14 +568,14 @@ class State(): report = [] # Source - report.append(ui.color_string('Source', 'GREEN')) + report.append(ansi.color_string('Source', 'GREEN')) report.extend(build_object_report(self.source)) report.append(' ') # Destination - report.append(ui.color_string('Destination', 'GREEN')) + report.append(ansi.color_string('Destination', 'GREEN')) if self.mode == 'Clone': - report[-1] += ui.color_string(' (ALL DATA WILL BE DELETED)', 'RED') + report[-1] += ansi.color_string(' (ALL DATA WILL BE DELETED)', 'RED') report.extend(build_object_report(self.destination)) report.append(' ') @@ -583,12 +583,12 @@ class State(): # NOTE: The check for block_pairs is to limit this section # to the second confirmation if self.mode == 'Clone' and self.block_pairs: - report.append(ui.color_string('WARNING', 'YELLOW')) + report.append(ansi.color_string('WARNING', 'YELLOW')) report.append( 'All data will be deleted from the destination listed above.', ) report.append( - ui.color_string( + ansi.color_string( ['This is irreversible and will lead to', 'DATA LOSS.'], ['YELLOW', 'RED'], ), @@ -607,18 +607,18 @@ class State(): # Map dir if self.working_dir: - report.append(ui.color_string('Map Save Directory', 'GREEN')) + report.append(ansi.color_string('Map Save Directory', 'GREEN')) report.append(f'{self.working_dir}/') report.append(' ') if not fstype_is_ok(self.working_dir, map_dir=True): report.append( - ui.color_string( + ansi.color_string( 'Map file(s) are being saved to a non-recommended filesystem.', 'YELLOW', ), ) report.append( - ui.color_string( + ansi.color_string( ['This is strongly discouraged and may lead to', 'DATA LOSS'], [None, 'RED'], ), @@ -627,11 +627,11 @@ class State(): # Source part(s) selected if source_parts: - report.append(ui.color_string('Source Part(s) selected', 'GREEN')) + report.append(ansi.color_string('Source Part(s) selected', 'GREEN')) if self.source.path.samefile(source_parts[0].path): report.append('Whole Disk') else: - report.append(ui.color_string(f'{"NAME":<9} SIZE', 'BLUE')) + report.append(ansi.color_string(f'{"NAME":<9} SIZE', 'BLUE')) for part in source_parts: report.append( f'{part.path.name:<9} ' @@ -663,7 +663,7 @@ class State(): error_size = self.get_error_size() error_size_str = std.bytes_to_string(error_size, decimals=2) if error_size > 0: - error_size_str = ui.color_string(error_size_str, 'YELLOW') + error_size_str = ansi.color_string(error_size_str, 'YELLOW') percent = self.get_percent_recovered() percent = format_status_string(percent, width=0) report.append(f'Overall rescued: {percent}, error size: {error_size_str}') @@ -675,7 +675,7 @@ class State(): error_size = pair.get_error_size() error_size_str = std.bytes_to_string(error_size, decimals=2) if error_size > 0: - error_size_str = ui.color_string(error_size_str, 'YELLOW') + error_size_str = ansi.color_string(error_size_str, 'YELLOW') pair_size = std.bytes_to_string(pair.size, decimals=2) percent = pair.get_percent_recovered() percent = format_status_string(percent, width=0) @@ -1066,10 +1066,10 @@ class State(): width = cfg.ddrescue.TMUX_SIDE_WIDTH # Status - report.append(ui.color_string(f'{"Status":^{width}}', 'BLUE')) + report.append(ansi.color_string(f'{"Status":^{width}}', 'BLUE')) if 'NEEDS ATTENTION' in overall_status: report.append( - ui.color_string(f'{overall_status:^{width}}', 'YELLOW_BLINK'), + ansi.color_string(f'{overall_status:^{width}}', 'YELLOW_BLINK'), ) else: report.append(f'{overall_status:^{width}}') @@ -1079,12 +1079,12 @@ class State(): if self.block_pairs: total_rescued = self.get_rescued_size() percent = self.get_percent_recovered() - report.append(ui.color_string('Overall Progress', 'BLUE')) + report.append(ansi.color_string('Overall Progress', 'BLUE')) report.append( f'Rescued: {format_status_string(percent, width=width-9)}', ) report.append( - ui.color_string( + ansi.color_string( [f'{std.bytes_to_string(total_rescued, decimals=2):>{width}}'], [get_percent_color(percent)], ), @@ -1093,7 +1093,7 @@ class State(): # Block pair progress for pair in self.block_pairs: - report.append(ui.color_string(pair.source, 'BLUE')) + report.append(ansi.color_string(pair.source, 'BLUE')) for name, status in pair.status.items(): name = name.title() report.append( @@ -1105,9 +1105,9 @@ class State(): if overall_status in ('Active', 'NEEDS ATTENTION'): etoc = get_etoc() report.append(separator) - report.append(ui.color_string('Estimated Pass Finish', 'BLUE')) + report.append(ansi.color_string('Estimated Pass Finish', 'BLUE')) if overall_status == 'NEEDS ATTENTION' or etoc == 'N/A': - report.append(ui.color_string('N/A', 'YELLOW')) + report.append(ansi.color_string('N/A', 'YELLOW')) else: report.append(etoc) @@ -1168,7 +1168,7 @@ class State(): source_str = _format_string(self.source, width) tmux.respawn_pane( self.panes['Source'], - text=ui.color_string( + text=ansi.color_string( ['Source', '' if source_exists else ' (Missing)', '\n', source_str], ['BLUE', 'RED', None, None], sep='', @@ -1183,7 +1183,7 @@ class State(): percent=50, vertical=False, target_id=self.panes['Source'], - text=ui.color_string( + text=ansi.color_string( ['Destination', '' if dest_exists else ' (Missing)', '\n', dest_str], ['BLUE', 'RED', None, None], sep='', @@ -1197,7 +1197,7 @@ def build_block_pair_report(block_pairs, settings): report = [] notes = [] if block_pairs: - report.append(ui.color_string('Block Pairs', 'GREEN')) + report.append(ansi.color_string('Block Pairs', 'GREEN')) else: # Bail early return report @@ -1216,7 +1216,7 @@ def build_block_pair_report(block_pairs, settings): if settings: if not settings['First Run']: notes.append( - ui.color_string( + ansi.color_string( ['NOTE:', 'Clone settings loaded from previous run.'], ['BLUE', None], ), @@ -1224,14 +1224,14 @@ def build_block_pair_report(block_pairs, settings): if settings['Needs Format'] and settings['Table Type']: msg = f'Destination will be formatted using {settings["Table Type"]}' notes.append( - ui.color_string( + ansi.color_string( ['NOTE:', msg], ['BLUE', None], ), ) if any(pair.get_rescued_size() > 0 for pair in block_pairs): notes.append( - ui.color_string( + ansi.color_string( ['NOTE:', 'Resume data loaded from map file(s).'], ['BLUE', None], ), @@ -1313,12 +1313,12 @@ def build_directory_report(path): for line in proc.stdout.splitlines(): line = line.replace('\n', '') if 'FSTYPE' in line: - line = ui.color_string(f'{"PATH":<{width}}{line}', 'BLUE') + line = ansi.color_string(f'{"PATH":<{width}}{line}', 'BLUE') else: line = f'{path:<{width}}{line}' report.append(line) else: - report.append(ui.color_string('PATH', 'BLUE')) + report.append(ansi.color_string('PATH', 'BLUE')) report.append(str(path)) # Done @@ -1354,7 +1354,7 @@ def build_disk_report(dev): # Partition details report.append( - ui.color_string( + ansi.color_string( ( f'{"NAME":<{widths["name"]}}' f'{" " if dev.children else ""}' @@ -1400,7 +1400,7 @@ def build_disk_report(dev): def build_main_menu(): """Build main menu, returns wk.ui.cli.Menu.""" - menu = ui.Menu(title=ui.color_string('ddrescue TUI: Main Menu', 'GREEN')) + menu = ui.Menu(title=ansi.color_string('ddrescue TUI: Main Menu', 'GREEN')) menu.separator = ' ' # Add actions, options, etc @@ -1433,9 +1433,9 @@ def build_object_report(obj): def build_settings_menu(silent=True): """Build settings menu, returns wk.ui.cli.Menu.""" title_text = [ - ui.color_string('ddrescue TUI: Expert Settings', 'GREEN'), + ansi.color_string('ddrescue TUI: Expert Settings', 'GREEN'), ' ', - ui.color_string( + ansi.color_string( ['These settings can cause', 'MAJOR DAMAGE', 'to drives'], ['YELLOW', 'RED', 'YELLOW'], ), @@ -1580,7 +1580,7 @@ def format_status_string(status, width): # Add color if necessary if color: - status_str = ui.color_string(status_str, color) + status_str = ansi.color_string(status_str, color) # Done return status_str @@ -1960,7 +1960,7 @@ def main(): # Save results to log LOG.info('') for line in state.generate_report(): - LOG.info(' %s', ui.strip_colors(line)) + LOG.info(' %s', ansi.strip_colors(line)) def mount_raw_image(path): @@ -2084,7 +2084,7 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): now = datetime.datetime.now(tz=TIMEZONE).strftime('%Y-%m-%d %H:%M %Z') with open(f'{state.log_dir}/smart.out', 'w', encoding='utf-8') as _f: _f.write( - ui.color_string( + ansi.color_string( ['SMART Attributes', f'Updated: {now}\n'], ['BLUE', 'YELLOW'], sep='\t\t', @@ -2267,7 +2267,7 @@ def select_disk(prompt, skip_disk=None): ui.print_info('Scanning disks...') disks = hw_disk.get_disks() menu = ui.Menu( - title=ui.color_string(f'ddrescue TUI: {prompt} Selection', 'GREEN'), + title=ansi.color_string(f'ddrescue TUI: {prompt} Selection', 'GREEN'), ) menu.disabled_str = 'Already selected' menu.separator = ' ' @@ -2309,7 +2309,7 @@ def select_disk(prompt, skip_disk=None): def select_disk_parts(prompt, disk): """Select disk parts from list, returns list of Disk().""" - title = ui.color_string('ddrescue TUI: Partition Selection', 'GREEN') + title = ansi.color_string('ddrescue TUI: Partition Selection', 'GREEN') title += f'\n\nDisk: {disk.path} {disk.description}' menu = ui.Menu(title) menu.separator = ' ' @@ -2360,7 +2360,7 @@ def select_disk_parts(prompt, disk): if not menu.options: menu.add_option(whole_disk_str, {'Selected': True, 'Path': disk.path}) menu.title += '\n\n' - menu.title += ui.color_string(' No partitions detected.', 'YELLOW') + menu.title += ansi.color_string(' No partitions detected.', 'YELLOW') # Get selection _select_parts(menu) @@ -2391,7 +2391,7 @@ def select_path(prompt): """Select path, returns pathlib.Path.""" invalid = False menu = ui.Menu( - title=ui.color_string(f'ddrescue TUI: {prompt} Path Selection', 'GREEN'), + title=ansi.color_string(f'ddrescue TUI: {prompt} Path Selection', 'GREEN'), ) menu.separator = ' ' menu.add_action('Quit') diff --git a/scripts/wk/graph.py b/scripts/wk/graph.py index 0de162d5..70e26cfb 100644 --- a/scripts/wk/graph.py +++ b/scripts/wk/graph.py @@ -3,7 +3,7 @@ import logging -from wk.ui import cli as ui +from wk import ansi # STATIC VARIABLES @@ -52,27 +52,27 @@ def generate_horizontal_graph(rate_list, graph_width=40, oneline=False): rate_color = 'GREEN' # Build graph - full_block = ui.color_string((GRAPH_HORIZONTAL[-1],), (rate_color,)) + full_block = ansi.color_string((GRAPH_HORIZONTAL[-1],), (rate_color,)) if step >= 24: - graph[0] += ui.color_string((GRAPH_HORIZONTAL[step-24],), (rate_color,)) + graph[0] += ansi.color_string((GRAPH_HORIZONTAL[step-24],), (rate_color,)) graph[1] += full_block graph[2] += full_block graph[3] += full_block elif step >= 16: graph[0] += ' ' - graph[1] += ui.color_string((GRAPH_HORIZONTAL[step-16],), (rate_color,)) + graph[1] += ansi.color_string((GRAPH_HORIZONTAL[step-16],), (rate_color,)) graph[2] += full_block graph[3] += full_block elif step >= 8: graph[0] += ' ' graph[1] += ' ' - graph[2] += ui.color_string((GRAPH_HORIZONTAL[step-8],), (rate_color,)) + graph[2] += ansi.color_string((GRAPH_HORIZONTAL[step-8],), (rate_color,)) graph[3] += full_block else: graph[0] += ' ' graph[1] += ' ' graph[2] += ' ' - graph[3] += ui.color_string((GRAPH_HORIZONTAL[step],), (rate_color,)) + graph[3] += ansi.color_string((GRAPH_HORIZONTAL[step],), (rate_color,)) # Done if oneline: @@ -128,7 +128,7 @@ def vertical_graph_line(percent, rate, scale=32): color_rate = 'GREEN' # Build string - line = ui.color_string( + line = ansi.color_string( strings=( f'{percent:5.1f}%', f'{GRAPH_VERTICAL[step]:<4}', diff --git a/scripts/wk/hw/benchmark.py b/scripts/wk/hw/benchmark.py index 58d3ee83..5a33ad88 100644 --- a/scripts/wk/hw/benchmark.py +++ b/scripts/wk/hw/benchmark.py @@ -5,7 +5,7 @@ import logging from subprocess import PIPE, STDOUT -from wk import graph +from wk import ansi, graph from wk.cfg.hw import ( IO_ALT_TEST_SIZE_FACTOR, IO_BLOCK_SIZE, @@ -22,7 +22,6 @@ from wk.cfg.hw import ( ) from wk.exe import run_program from wk.std import PLATFORM -from wk.ui import cli as ui # STATIC VARIABLES @@ -113,7 +112,7 @@ def check_io_results(test_obj, rate_list, graph_width) -> None: # Add horizontal graph to report for line in graph.generate_horizontal_graph(rate_list, graph_width): - if not ui.strip_colors(line).strip(): + if not ansi.strip_colors(line).strip(): # Skip empty lines continue test_obj.report.append(line) @@ -151,7 +150,7 @@ def run_io_test(test_obj, log_path, test_mode=False) -> None: LOG.info('Using %s for better performance', dev_path) offset = 0 read_rates = [] - test_obj.report.append(ui.color_string('I/O Benchmark', 'BLUE')) + test_obj.report.append(ansi.color_string('I/O Benchmark', 'BLUE')) # Get dd values or bail try: @@ -159,7 +158,7 @@ def run_io_test(test_obj, log_path, test_mode=False) -> None: except DeviceTooSmallError: test_obj.set_status('N/A') test_obj.report.append( - ui.color_string('Disk too small to test', 'YELLOW'), + ansi.color_string('Disk too small to test', 'YELLOW'), ) return diff --git a/scripts/wk/hw/cpu.py b/scripts/wk/hw/cpu.py index 9a741ecd..af8b0c46 100644 --- a/scripts/wk/hw/cpu.py +++ b/scripts/wk/hw/cpu.py @@ -7,7 +7,7 @@ import subprocess from typing import TextIO -from wk import exe +from wk import ansi, exe from wk.cfg.hw import CPU_FAILURE_TEMP from wk.os.mac import set_fans as macos_set_fans from wk.std import PLATFORM @@ -93,9 +93,9 @@ def check_mprime_results(test_obj, working_dir) -> None: for line in passing_lines: test_obj.report.append(f' {line}') for line in warning_lines: - test_obj.report.append(ui.color_string(f' {line}', 'YELLOW')) + test_obj.report.append(ansi.color_string(f' {line}', 'YELLOW')) if not (passing_lines or warning_lines): - test_obj.report.append(ui.color_string(' Unknown result', 'YELLOW')) + test_obj.report.append(ansi.color_string(' Unknown result', 'YELLOW')) def start_mprime(working_dir, log_path) -> subprocess.Popen: diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py index 72c594e4..0ad9206c 100644 --- a/scripts/wk/hw/diags.py +++ b/scripts/wk/hw/diags.py @@ -10,7 +10,7 @@ import time from docopt import docopt -from wk import cfg, debug, exe, log, std +from wk import ansi, cfg, debug, exe, log, std from wk.cfg.hw import STATUS_COLORS from wk.hw import benchmark as hw_benchmark from wk.hw import cpu as hw_cpu @@ -89,9 +89,9 @@ class State(): self.panes = {} self.system = None self.test_groups = [] - self.top_text = ui.color_string('Hardware Diagnostics', 'GREEN') + self.top_text = ansi.color_string('Hardware Diagnostics', 'GREEN') if test_mode: - self.top_text += ui.color_string(' (Test Mode)', 'YELLOW') + self.top_text += ansi.color_string(' (Test Mode)', 'YELLOW') # Init tmux and start a background process to maintain layout self.init_tmux() @@ -229,7 +229,7 @@ class State(): self.panes['Started'] = tmux.split_window( lines=cfg.hw.TMUX_SIDE_WIDTH, target_id=self.panes['Top'], - text=ui.color_string( + text=ansi.color_string( ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], ['BLUE', None], sep='\n', @@ -292,7 +292,7 @@ class State(): """Update 'Started' pane following clock sync.""" tmux.respawn_pane( pane_id=self.panes['Started'], - text=ui.color_string( + text=ansi.color_string( ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], ['BLUE', None], sep='\n', @@ -305,9 +305,9 @@ class State(): width = cfg.hw.TMUX_SIDE_WIDTH for group in self.test_groups: - report.append(ui.color_string(group.name, 'BLUE')) + report.append(ansi.color_string(group.name, 'BLUE')) for test in group.test_objects: - report.append(ui.color_string( + report.append(ansi.color_string( [test.label, f'{test.status:>{width-len(test.label)}}'], [None, STATUS_COLORS.get(test.status, None)], sep='', @@ -453,7 +453,7 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: sensors.save_average_temps(temp_label='Cooldown', seconds=5) # Check Prime95 results - test_mprime_obj.report.append(ui.color_string('Prime95', 'BLUE')) + test_mprime_obj.report.append(ansi.color_string('Prime95', 'BLUE')) hw_cpu.check_mprime_results( test_obj=test_mprime_obj, working_dir=state.log_dir, ) @@ -493,7 +493,7 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: state.update_progress_pane() # Check Cooling results - test_cooling_obj.report.append(ui.color_string('Temps', 'BLUE')) + test_cooling_obj.report.append(ansi.color_string('Temps', 'BLUE')) hw_cpu.check_cooling_results(test_cooling_obj, sensors, run_sysbench) # Cleanup @@ -566,12 +566,12 @@ def disk_io_benchmark( # Something went wrong LOG.error('%s', err) test.set_status('ERROR') - test.report.append(ui.color_string(' Unknown Error', 'RED')) + test.report.append(ansi.color_string(' Unknown Error', 'RED')) # Mark test(s) aborted if necessary if aborted: test.set_status('Aborted') - test.report.append(ui.color_string(' Aborted', 'YELLOW')) + test.report.append(ansi.color_string(' Aborted', 'YELLOW')) break # Update progress after each test @@ -734,7 +734,7 @@ def disk_surface_scan(state, test_objects, test_mode=False) -> None: for test in test_objects: if not (test.disabled or test.passed or test.failed): test.set_status('Aborted') - test.report.append(ui.color_string(' Aborted', 'YELLOW')) + test.report.append(ansi.color_string(' Aborted', 'YELLOW')) # Cleanup state.update_progress_pane() diff --git a/scripts/wk/hw/disk.py b/scripts/wk/hw/disk.py index 26a3d0c5..bae3e305 100644 --- a/scripts/wk/hw/disk.py +++ b/scripts/wk/hw/disk.py @@ -11,6 +11,7 @@ import re from dataclasses import dataclass, field from typing import Any, Union +from wk import ansi from wk.cfg.main import KIT_NAME_SHORT from wk.cfg.python import DATACLASS_DECORATOR_KWARGS from wk.exe import get_json_from_command, run_program @@ -20,7 +21,6 @@ from wk.hw.smart import ( get_known_disk_attributes, ) from wk.std import PLATFORM -from wk.ui import cli as ui # STATIC VARIABLES @@ -74,7 +74,7 @@ class Disk: def add_note(self, note, color=None) -> None: """Add note that will be included in the disk report.""" if color: - note = ui.color_string(note, color) + note = ansi.color_string(note, color) if note not in self.notes: self.notes.append(note) self.notes.sort() @@ -83,7 +83,7 @@ class Disk: """Check if note is already present.""" present = False for note in self.notes: - if note_str == ui.strip_colors(note): + if note_str == ansi.strip_colors(note): present = True return present @@ -99,18 +99,18 @@ class Disk: """Generate Disk report, returns list.""" report = [] if header: - report.append(ui.color_string(f'Device ({self.path.name})', 'BLUE')) + report.append(ansi.color_string(f'Device ({self.path.name})', 'BLUE')) report.append(f' {self.description}') # Attributes if self.attributes: if header: - report.append(ui.color_string('Attributes', 'BLUE')) + report.append(ansi.color_string('Attributes', 'BLUE')) report.extend(generate_attribute_report(self)) # Notes if self.notes: - report.append(ui.color_string('Notes', 'BLUE')) + report.append(ansi.color_string('Notes', 'BLUE')) for note in self.notes: report.append(f' {note}') diff --git a/scripts/wk/hw/sensors.py b/scripts/wk/hw/sensors.py index 858b6553..f05a628d 100644 --- a/scripts/wk/hw/sensors.py +++ b/scripts/wk/hw/sensors.py @@ -9,11 +9,11 @@ import re from subprocess import CalledProcessError from typing import Any +from wk import ansi from wk.cfg.hw import CPU_CRITICAL_TEMP, SMC_IDS, TEMP_COLORS from wk.exe import run_program, start_thread from wk.io import non_clobber_path from wk.std import PLATFORM, sleep -from wk.ui import cli as ui # STATIC VARIABLES @@ -110,7 +110,7 @@ class Sensors(): # Handle empty reports if not report: report = [ - ui.color_string('WARNING: No sensors found', 'YELLOW'), + ansi.color_string('WARNING: No sensors found', 'YELLOW'), '', 'Please monitor temps manually', ] @@ -426,7 +426,7 @@ def get_temp_str(temp, colored=True) -> str: temp = float(temp) except (TypeError, ValueError): # Invalid temp? - return ui.color_string(temp, 'PURPLE') + return ansi.color_string(temp, 'PURPLE') # Determine color if colored: @@ -436,7 +436,7 @@ def get_temp_str(temp, colored=True) -> str: break # Done - return ui.color_string(f'{"-" if temp < 0 else ""}{temp:2.0f}°C', temp_color) + return ansi.color_string(f'{"-" if temp < 0 else ""}{temp:2.0f}°C', temp_color) diff --git a/scripts/wk/hw/smart.py b/scripts/wk/hw/smart.py index c777f400..88c016ab 100644 --- a/scripts/wk/hw/smart.py +++ b/scripts/wk/hw/smart.py @@ -7,6 +7,7 @@ import re from typing import Any +from wk import ansi from wk.cfg.hw import ( ATTRIBUTE_COLORS, KEY_NVME, @@ -19,7 +20,6 @@ from wk.cfg.hw import ( ) from wk.exe import get_json_from_command, run_program from wk.std import bytes_to_string, sleep -from wk.ui import cli as ui # STATIC VARIABLES @@ -41,26 +41,26 @@ def build_self_test_report(test_obj, aborted=False) -> None: For instance if the test was aborted the report should include the last known progress instead of just "was aborted by host." """ - report = [ui.color_string('Self-Test', 'BLUE')] + report = [ansi.color_string('Self-Test', 'BLUE')] test_details = get_smart_self_test_details(test_obj.dev) test_result = test_details.get('status', {}).get('string', 'Unknown') # Build report if test_obj.disabled or test_obj.status == 'Denied': - report.append(ui.color_string(f' {test_obj.status}', 'RED')) + report.append(ansi.color_string(f' {test_obj.status}', 'RED')) elif test_obj.status == 'N/A' or not test_obj.dev.attributes: - report.append(ui.color_string(f' {test_obj.status}', 'YELLOW')) + report.append(ansi.color_string(f' {test_obj.status}', 'YELLOW')) elif test_obj.status == 'TestInProgress': - report.append(ui.color_string(' Failed to stop previous test', 'RED')) + report.append(ansi.color_string(' Failed to stop previous test', 'RED')) test_obj.set_status('Failed') else: # Other cases include self-test result string report.append(f' {test_result.capitalize()}') if aborted and not (test_obj.passed or test_obj.failed): - report.append(ui.color_string(' Aborted', 'YELLOW')) + report.append(ansi.color_string(' Aborted', 'YELLOW')) test_obj.set_status('Aborted') elif test_obj.status == 'TimedOut': - report.append(ui.color_string(' TimedOut', 'YELLOW')) + report.append(ansi.color_string(' TimedOut', 'YELLOW')) # Done test_obj.report.extend(report) @@ -137,7 +137,7 @@ def generate_attribute_report(dev, only_failed=False) -> list[str]: continue # Build colored string and append to report - line = ui.color_string( + line = ansi.color_string( [label, get_attribute_value_string(dev, attr), note], [None, value_color, 'YELLOW'], ) @@ -299,7 +299,7 @@ def run_smart_self_test(test_obj, log_path) -> bool: finished = False test_details = get_smart_self_test_details(test_obj.dev) size_str = bytes_to_string(test_obj.dev.size, use_binary=False) - header_str = ui.color_string( + header_str = ansi.color_string( ['[', test_obj.dev.path.name, ' ', size_str, ']'], [None, 'BLUE', None, 'CYAN', None], sep='', diff --git a/scripts/wk/hw/surface_scan.py b/scripts/wk/hw/surface_scan.py index 9a542d30..7f4ed243 100644 --- a/scripts/wk/hw/surface_scan.py +++ b/scripts/wk/hw/surface_scan.py @@ -5,6 +5,7 @@ import logging from subprocess import STDOUT +from wk import ansi from wk.cfg.hw import ( BADBLOCKS_EXTRA_LARGE_DISK, BADBLOCKS_LARGE_DISK, @@ -15,7 +16,6 @@ from wk.cfg.hw import ( ) from wk.exe import run_program from wk.std import PLATFORM, bytes_to_string -from wk.ui import cli as ui # STATIC VARIABLES @@ -27,7 +27,7 @@ def check_surface_scan_results(test_obj, log_path) -> None: """Check results and set test status.""" with open(log_path, 'r', encoding='utf-8') as _f: for line in _f.readlines(): - line = ui.strip_colors(line.strip()) + line = ansi.strip_colors(line.strip()) if not line or BADBLOCKS_SKIP_REGEX.match(line): # Skip continue @@ -44,10 +44,10 @@ def check_surface_scan_results(test_obj, log_path) -> None: test_obj.set_status('Passed') else: test_obj.failed = True - test_obj.report.append(f' {ui.color_string(line, "YELLOW")}') + test_obj.report.append(f' {ansi.color_string(line, "YELLOW")}') test_obj.set_status('Failed') else: - test_obj.report.append(f' {ui.color_string(line, "YELLOW")}') + test_obj.report.append(f' {ansi.color_string(line, "YELLOW")}') if not (test_obj.passed or test_obj.failed): test_obj.set_status('Unknown') @@ -61,7 +61,7 @@ def run_scan(test_obj, log_path, test_mode=False) -> None: # Use "RAW" disks under macOS dev_path = dev_path.with_name(f'r{dev_path.name}') LOG.info('Using %s for better performance', dev_path) - test_obj.report.append(ui.color_string('badblocks', 'BLUE')) + test_obj.report.append(ansi.color_string('badblocks', 'BLUE')) test_obj.set_status('Working') # Increase block size if necessary @@ -80,7 +80,7 @@ def run_scan(test_obj, log_path, test_mode=False) -> None: with open(log_path, 'a', encoding='utf-8') as _f: size_str = bytes_to_string(dev.size, use_binary=False) _f.write( - ui.color_string( + ansi.color_string( ['[', dev.path.name, ' ', size_str, ']\n'], [None, 'BLUE', None, 'CYAN', None], sep='', diff --git a/scripts/wk/hw/system.py b/scripts/wk/hw/system.py index b82f9433..949adb38 100644 --- a/scripts/wk/hw/system.py +++ b/scripts/wk/hw/system.py @@ -8,12 +8,12 @@ import re from dataclasses import dataclass, field from typing import Any +from wk import ansi from wk.cfg.hw import KNOWN_RAM_VENDOR_IDS from wk.cfg.python import DATACLASS_DECORATOR_KWARGS from wk.exe import get_json_from_command, run_program from wk.hw.test import Test from wk.std import PLATFORM, bytes_to_string, string_to_bytes -from wk.ui import cli as ui # STATIC VARIABLES @@ -37,11 +37,11 @@ class System: def generate_report(self) -> list[str]: """Generate CPU & RAM report, returns list.""" report = [] - report.append(ui.color_string('Device', 'BLUE')) + report.append(ansi.color_string('Device', 'BLUE')) report.append(f' {self.cpu_description}') # Include RAM details - report.append(ui.color_string('RAM', 'BLUE')) + report.append(ansi.color_string('RAM', 'BLUE')) report.append(f' {self.ram_total} ({", ".join(self.ram_dimms)})') # Tests diff --git a/scripts/wk/os/linux.py b/scripts/wk/os/linux.py index 991dd454..a4c82002 100644 --- a/scripts/wk/os/linux.py +++ b/scripts/wk/os/linux.py @@ -7,11 +7,11 @@ import pathlib import re import subprocess +from wk import ansi from wk.cfg.hw import VOLUME_FAILURE_THRESHOLD, VOLUME_WARNING_THRESHOLD from wk.exe import get_json_from_command, popen_program, run_program from wk.log import format_log_path from wk.std import bytes_to_string -from wk.ui import cli as ui # STATIC VARIABLES @@ -83,20 +83,20 @@ def build_volume_report(device_path=None) -> list: vol['mountpoint'] = f'Mounted on {vol["mountpoint"]}' # Name and size - line = ui.color_string( + line = ansi.color_string( [f'{vol["name"]:<20}', f'{vol["size"]:>9}'], [None, 'CYAN'], ) # Mountpoint and type - line = ui.color_string( + line = ansi.color_string( [line, f'{vol["mountpoint"]:<{m_width}}', f'{vol["fstype"]:<11}'], [None, None, 'BLUE'], ) # Used and free if any([vol['fsused'], vol['fsavail']]): - line = ui.color_string( + line = ansi.color_string( [line, f'({vol["fsused"]:>9} used, {vol["fsavail"]:>9} free)'], [None, size_color], ) diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py index 852f6065..3b2b9d93 100644 --- a/scripts/wk/os/win.py +++ b/scripts/wk/os/win.py @@ -17,6 +17,7 @@ except ImportError as err: if platform.system() == 'Windows': raise err +from wk import ansi from wk.borrowed import acpi from wk.cfg.main import KIT_NAME_FULL from wk.cfg.windows_builds import ( @@ -182,7 +183,7 @@ def check_4k_alignment(show_alert=False): continue if int(match.group('offset')) % 4096 != 0: report.append( - ui.color_string( + ansi.color_string( f'{match.group("description")}' f' ({bytes_to_string(match.group("size"), decimals=1)})' , @@ -198,7 +199,7 @@ def check_4k_alignment(show_alert=False): if report: report.insert( 0, - ui.color_string('One or more partitions not 4K aligned', 'YELLOW'), + ansi.color_string('One or more partitions not 4K aligned', 'YELLOW'), ) return report @@ -250,13 +251,13 @@ def get_installed_antivirus(): state = proc.stdout.split('=')[1] state = hex(int(state)) if str(state)[3:5] not in ['10', '11']: - report.append(ui.color_string(f'[Disabled] {product}', 'YELLOW')) + report.append(ansi.color_string(f'[Disabled] {product}', 'YELLOW')) else: report.append(product) # Final check if not report: - report.append(ui.color_string('No products detected', 'RED')) + report.append(ansi.color_string('No products detected', 'RED')) # Done return report @@ -363,7 +364,7 @@ def get_volume_usage(use_colors=False): f' ({bytes_to_string(free, 2):>10} / {bytes_to_string(total, 2):>10})' ) if use_colors: - display_str = ui.color_string(display_str, color) + display_str = ansi.color_string(display_str, color) report.append(f'{disk.device} {display_str}') # Done diff --git a/scripts/wk/repairs/win.py b/scripts/wk/repairs/win.py index d94f6e3d..5ae419ed 100644 --- a/scripts/wk/repairs/win.py +++ b/scripts/wk/repairs/win.py @@ -11,6 +11,7 @@ import time from subprocess import CalledProcessError, DEVNULL from xml.dom.minidom import parse as xml_parse +from wk import ansi from wk.cfg.main import KIT_NAME_FULL, KIT_NAME_SHORT, WINDOWS_TIME_ZONE from wk.cfg.repairs import ( AUTO_REPAIR_DELAY_IN_SECONDS, @@ -104,7 +105,7 @@ for error in ('CalledProcessError', 'FileNotFoundError'): def build_menus(base_menus, title, presets): """Build menus, returns dict.""" menus = {} - menus['Main'] = ui.Menu(title=f'{title}\n{ui.color_string("Main Menu", "GREEN")}') + menus['Main'] = ui.Menu(title=f'{title}\n{ansi.color_string("Main Menu", "GREEN")}') # Main Menu for entry in base_menus['Actions']: @@ -113,7 +114,7 @@ def build_menus(base_menus, title, presets): menus['Main'].add_option(group, {'Selected': True}) # Options - menus['Options'] = ui.Menu(title=f'{title}\n{ui.color_string("Options", "GREEN")}') + menus['Options'] = ui.Menu(title=f'{title}\n{ansi.color_string("Options", "GREEN")}') for entry in base_menus['Options']: menus['Options'].add_option(entry.name, entry.details) menus['Options'].add_action('All') @@ -123,7 +124,7 @@ def build_menus(base_menus, title, presets): # Run groups for group, entries in base_menus['Groups'].items(): - menus[group] = ui.Menu(title=f'{title}\n{ui.color_string(group, "GREEN")}') + menus[group] = ui.Menu(title=f'{title}\n{ansi.color_string(group, "GREEN")}') menus[group].disabled_str = 'Locked' for entry in entries: menus[group].add_option(entry.name, entry.details) @@ -155,7 +156,7 @@ def build_menus(base_menus, title, presets): ) # Update presets Menu - MENU_PRESETS.title = f'{title}\n{ui.color_string("Load Preset", "GREEN")}' + MENU_PRESETS.title = f'{title}\n{ansi.color_string("Load Preset", "GREEN")}' MENU_PRESETS.add_option('Default') for name in presets: MENU_PRESETS.add_option(name) @@ -380,7 +381,7 @@ def load_settings(menus): if group == 'Main': continue for name in menu.options: - menu.options[name].update(get_entry_settings(group, ui.strip_colors(name))) + menu.options[name].update(get_entry_settings(group, ansi.strip_colors(name))) def run_auto_repairs(base_menus, presets): @@ -446,7 +447,7 @@ def run_group(group, menu): """Run entries in group if appropriate.""" ui.print_info(f' {group}') for name, details in menu.options.items(): - name_str = ui.strip_colors(name) + name_str = ansi.strip_colors(name) skipped = details.get('Skipped', False) done = details.get('Done', False) disabled = details.get('Disabled', False) @@ -501,7 +502,7 @@ def save_selection_settings(menus): def save_settings(group, name, result=None, **kwargs): """Save entry settings in the registry.""" - key_path = fr'{AUTO_REPAIR_KEY}\{group}\{ui.strip_colors(name)}' + key_path = fr'{AUTO_REPAIR_KEY}\{group}\{ansi.strip_colors(name)}' # Get values from TryAndPrint result if result: @@ -515,7 +516,7 @@ def save_settings(group, name, result=None, **kwargs): # Write values to registry for value_name, data in kwargs.items(): - value_name = ui.strip_colors(value_name) + value_name = ansi.strip_colors(value_name) if isinstance(data, bool): data = 1 if data else 0 if isinstance(data, int): diff --git a/scripts/wk/setup/win.py b/scripts/wk/setup/win.py index 22d8224f..96157af3 100644 --- a/scripts/wk/setup/win.py +++ b/scripts/wk/setup/win.py @@ -8,6 +8,7 @@ import os import re import sys +from wk import ansi from wk.cfg.main import KIT_NAME_FULL from wk.cfg.setup import ( BROWSER_PATHS, @@ -101,7 +102,7 @@ for error in ('CalledProcessError', 'FileNotFoundError'): def build_menus(base_menus, title, presets): """Build menus, returns dict.""" menus = {} - menus['Main'] = ui.Menu(title=f'{title}\n{ui.color_string("Main Menu", "GREEN")}') + menus['Main'] = ui.Menu(title=f'{title}\n{ansi.color_string("Main Menu", "GREEN")}') # Main Menu for entry in base_menus['Actions']: @@ -111,7 +112,7 @@ def build_menus(base_menus, title, presets): # Run groups for group, entries in base_menus['Groups'].items(): - menus[group] = ui.Menu(title=f'{title}\n{ui.color_string(group, "GREEN")}') + menus[group] = ui.Menu(title=f'{title}\n{ansi.color_string(group, "GREEN")}') for entry in entries: menus[group].add_option(entry.name, entry.details) menus[group].add_action('All') @@ -140,7 +141,7 @@ def build_menus(base_menus, title, presets): ) # Update presets Menu - MENU_PRESETS.title = f'{title}\n{ui.color_string("Load Preset", "GREEN")}' + MENU_PRESETS.title = f'{title}\n{ansi.color_string("Load Preset", "GREEN")}' MENU_PRESETS.add_option('Default') for name in presets: MENU_PRESETS.add_option(name) @@ -176,7 +177,7 @@ def check_os_and_set_menu_title(title): color = 'RED' # Done - return f'{title} ({ui.color_string(os_name, color)})' + return f'{title} ({ansi.color_string(os_name, color)})' def load_preset(menus, presets, title, enable_menu_exit=True): @@ -264,7 +265,7 @@ def run_group(group, menu): """Run entries in group if appropriate.""" ui.print_info(f' {group}') for name, details in menu.options.items(): - name_str = ui.strip_colors(name) + name_str = ansi.strip_colors(name) # Not selected if not details.get('Selected', False): @@ -893,7 +894,7 @@ def get_storage_status(): """Get storage status for fixed disks, returns list.""" report = get_volume_usage(use_colors=True) for disk in get_raw_disks(): - report.append(ui.color_string(f'Uninitialized Disk: {disk}', 'RED')) + report.append(ansi.color_string(f'Uninitialized Disk: {disk}', 'RED')) # Done return report diff --git a/scripts/wk/ui/cli.py b/scripts/wk/ui/cli.py index ecd8965a..6aa5a128 100644 --- a/scripts/wk/ui/cli.py +++ b/scripts/wk/ui/cli.py @@ -1,10 +1,8 @@ """WizardKit: CLI functions""" # vim: sts=2 sw=2 ts=2 -import itertools import logging import os -import pathlib import platform import re import subprocess @@ -19,6 +17,7 @@ except ImportError: # Assuming Python is < 3.9 from functools import lru_cache as cache +from wk.ansi import color_string, strip_colors from wk.cfg.main import ( ENABLED_UPLOAD_DATA, INDENT, @@ -28,19 +27,6 @@ from wk.cfg.main import ( from wk.std import (sleep, GenericWarning) # STATIC VARIABLES -COLORS = { - 'CLEAR': '\033[0m', - 'RED': '\033[31m', - 'RED_BLINK': '\033[31;5m', - 'ORANGE': '\033[31;1m', - 'ORANGE_RED': '\033[1;31;41m', - 'GREEN': '\033[32m', - 'YELLOW': '\033[33m', - 'YELLOW_BLINK': '\033[33;5m', - 'BLUE': '\033[34m', - 'PURPLE': '\033[35m', - 'CYAN': '\033[36m', - } LOG = logging.getLogger(__name__) PLATFORM = platform.system() @@ -646,38 +632,6 @@ def clear_screen(): print('\033c') -def color_string(strings, colors, sep=' '): - """Build colored string using ANSI escapes, returns str.""" - clear_code = COLORS['CLEAR'] - msg = [] - - # Convert to tuples if necessary - if isinstance(strings, (str, pathlib.Path)): - strings = (strings,) - if isinstance(colors, (str, pathlib.Path)): - colors = (colors,) - - # Convert to strings if necessary - try: - iter(strings) - except TypeError: - # Assuming single element passed, convert to string - strings = (str(strings),) - try: - iter(colors) - except TypeError: - # Assuming single element passed, convert to string - colors = (str(colors),) - - # Build new string with color escapes added - for string, color in itertools.zip_longest(strings, colors): - color_code = COLORS.get(color, clear_code) - msg.append(f'{color_code}{string}{clear_code}') - - # Done - return sep.join(msg) - - @cache def get_exception(name): """Get exception by name, returns exception object. @@ -850,13 +804,5 @@ def show_data(message, data, color=None, indent=None, width=None): ) -def strip_colors(string): - """Strip known ANSI color escapes from string, returns str.""" - LOG.debug('string: %s', string) - for color in COLORS.values(): - string = string.replace(color, '') - return string - - if __name__ == '__main__': print("This file is not meant to be called directly.") From 44ddb3c258fb4728575b8d0dbc663bcb091d08d4 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sat, 8 Apr 2023 12:45:45 -0700 Subject: [PATCH 05/16] Replace clear_screen() function This is now part of wk.ansi and is fully cross-compatible. --- scripts/wk/ansi.py | 5 +++++ scripts/wk/ui/cli.py | 12 +----------- 2 files changed, 6 insertions(+), 11 deletions(-) diff --git a/scripts/wk/ansi.py b/scripts/wk/ansi.py index 2acf515d..2c21c0ed 100644 --- a/scripts/wk/ansi.py +++ b/scripts/wk/ansi.py @@ -23,6 +23,11 @@ COLORS = { # Functions +def clear_screen(): + """Clear screen using ANSI escape.""" + print('\033c', end='', flush=True) + + def color_string(strings, colors, sep=' '): """Build colored string using ANSI escapes, returns str.""" clear_code = COLORS['CLEAR'] diff --git a/scripts/wk/ui/cli.py b/scripts/wk/ui/cli.py index 6aa5a128..298fecaf 100644 --- a/scripts/wk/ui/cli.py +++ b/scripts/wk/ui/cli.py @@ -17,7 +17,7 @@ except ImportError: # Assuming Python is < 3.9 from functools import lru_cache as cache -from wk.ansi import color_string, strip_colors +from wk.ansi import clear_screen, color_string, strip_colors from wk.cfg.main import ( ENABLED_UPLOAD_DATA, INDENT, @@ -622,16 +622,6 @@ def choice(choices, prompt='答えろ！'): return answer -def clear_screen(): - """Simple wrapper for clear/cls.""" - cmd = 'cls' if os.name == 'nt' else 'clear' - proc = subprocess.run(cmd, check=False, shell=True, stderr=subprocess.PIPE) - - # Workaround for live macOS env - if proc.returncode != 0: - print('\033c') - - @cache def get_exception(name): """Get exception by name, returns exception object. From 13fc64e6ab192a9a04316d44054656e03fa55b72 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sat, 8 Apr 2023 14:16:41 -0700 Subject: [PATCH 06/16] Remove unneeded wk.ui calls --- scripts/wk/hw/cpu.py | 9 ++++----- scripts/wk/hw/keyboard.py | 4 ++-- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/scripts/wk/hw/cpu.py b/scripts/wk/hw/cpu.py index af8b0c46..deb8f808 100644 --- a/scripts/wk/hw/cpu.py +++ b/scripts/wk/hw/cpu.py @@ -11,7 +11,6 @@ from wk import ansi, exe from wk.cfg.hw import CPU_FAILURE_TEMP from wk.os.mac import set_fans as macos_set_fans from wk.std import PLATFORM -from wk.ui import cli as ui from wk.ui import tmux @@ -66,7 +65,7 @@ def check_mprime_results(test_obj, working_dir) -> None: if re.search(r'(error|fail)', line, re.IGNORECASE): warning_lines[line] = None - # print.log (check if passed) + # prime.log (check if passed) for line in _read_file('prime.log'): line = line.strip() match = re.search( @@ -170,9 +169,9 @@ def set_apple_fan_speed(speed) -> None: except (RuntimeError, ValueError, subprocess.CalledProcessError) as err: LOG.error('Failed to set fans to %s', speed) LOG.error('Error: %s', err) - ui.print_error(f'Failed to set fans to {speed}') - for line in str(err).splitlines(): - ui.print_warning(f' {line.strip()}') + #ui.print_error(f'Failed to set fans to {speed}') + #for line in str(err).splitlines(): + # ui.print_warning(f' {line.strip()}') elif PLATFORM == 'Linux': cmd = ['apple-fans', speed] exe.run_program(cmd, check=False) diff --git a/scripts/wk/hw/keyboard.py b/scripts/wk/hw/keyboard.py index 83541c45..07d3a22d 100644 --- a/scripts/wk/hw/keyboard.py +++ b/scripts/wk/hw/keyboard.py @@ -5,7 +5,6 @@ import logging from wk.exe import run_program from wk.std import PLATFORM -from wk.ui import cli as ui # STATIC VARIABLES @@ -18,7 +17,8 @@ def keyboard_test() -> None: if PLATFORM == 'Linux': run_xev() else: - ui.print_warning(f'Not supported under this OS: {PLATFORM}') + LOG.error('Not supported under this OS: %s', PLATFORM) + raise NotImplementedError(f'Not supported under this OS: {PLATFORM}') def run_xev() -> None: From d302be2d7cf1c89ef1e12ea3efbddcafdc402feb Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sat, 8 Apr 2023 16:26:51 -0700 Subject: [PATCH 07/16] Use prompt_toolkit for CLI input --- scripts/wk/clone/ddrescue.py | 38 +++---- scripts/wk/os/win.py | 2 +- scripts/wk/ui/cli.py | 200 ++++++++++++++++++++++++----------- setup/windows/build.ps1 | 4 +- 4 files changed, 158 insertions(+), 86 deletions(-) diff --git a/scripts/wk/clone/ddrescue.py b/scripts/wk/clone/ddrescue.py index 672c3e25..2c8bd568 100644 --- a/scripts/wk/clone/ddrescue.py +++ b/scripts/wk/clone/ddrescue.py @@ -525,8 +525,8 @@ class State(): settings['Needs Format'] = True offset = 0 user_choice = ui.choice( - ['G', 'M', 'S'], 'Format clone using GPT, MBR, or match Source type?', + ['G', 'M', 'S'], ) if user_choice == 'G': settings['Table Type'] = 'GPT' @@ -563,7 +563,7 @@ class State(): bp_dest = self.destination self._add_block_pair(part, bp_dest) - def confirm_selections(self, prompt, source_parts=None): + def confirm_selections(self, prompt_msg, source_parts=None): """Show selection details and prompt for confirmation.""" report = [] @@ -642,7 +642,7 @@ class State(): # Prompt user ui.clear_screen() ui.print_report(report) - if not ui.ask(prompt): + if not ui.ask(prompt_msg): raise std.GenericAbort() def generate_report(self): @@ -746,7 +746,7 @@ class State(): # Confirmation #1 self.confirm_selections( - prompt='Are these selections correct?', + prompt_msg='Are these selections correct?', source_parts=source_parts, ) @@ -1449,7 +1449,7 @@ def build_settings_menu(silent=True): ui.print_standard( f'Available ddrescue presets: {" / ".join(SETTING_PRESETS)}' ) - preset = ui.choice(SETTING_PRESETS, 'Please select a preset:') + preset = ui.choice('Please select a preset:', SETTING_PRESETS) # Fix selection for _p in SETTING_PRESETS: @@ -1788,19 +1788,9 @@ def get_table_type(disk_path): def get_working_dir(mode, destination, force_local=False): """Get working directory using mode and destination, returns path.""" - ticket_id = None + ticket_id = ui.get_ticket_id() working_dir = None - # Set ticket ID - while ticket_id is None: - ticket_id = ui.input_text( - prompt='Please enter ticket ID:', - allow_empty_response=False, - ) - ticket_id = ticket_id.replace(' ', '_') - if not re.match(r'^\d+', ticket_id): - ticket_id = None - # Use preferred path if possible if mode == 'Image': try: @@ -2262,12 +2252,12 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True): state.update_progress_pane('Idle') -def select_disk(prompt, skip_disk=None): +def select_disk(prompt_msg, skip_disk=None): """Select disk from list, returns Disk().""" ui.print_info('Scanning disks...') disks = hw_disk.get_disks() menu = ui.Menu( - title=ansi.color_string(f'ddrescue TUI: {prompt} Selection', 'GREEN'), + title=ansi.color_string(f'ddrescue TUI: {prompt_msg} Selection', 'GREEN'), ) menu.disabled_str = 'Already selected' menu.separator = ' ' @@ -2307,7 +2297,7 @@ def select_disk(prompt, skip_disk=None): return selected_disk -def select_disk_parts(prompt, disk): +def select_disk_parts(prompt_msg, disk): """Select disk parts from list, returns list of Disk().""" title = ansi.color_string('ddrescue TUI: Partition Selection', 'GREEN') title += f'\n\nDisk: {disk.path} {disk.description}' @@ -2323,7 +2313,7 @@ def select_disk_parts(prompt, disk): """Loop over selection menu until at least one partition selected.""" while True: selection = menu.advanced_select( - f'Please select the parts to {prompt.lower()}: ', + f'Please select the parts to {prompt_msg.lower()}: ', ) if 'All' in selection: for option in menu.options.values(): @@ -2373,7 +2363,7 @@ def select_disk_parts(prompt, disk): # Check if whole disk selected if len(object_list) == len(disk.children): # NOTE: This is not true if the disk has no partitions - msg = f'Preserve partition table and unused space in {prompt.lower()}?' + msg = f'Preserve partition table and unused space in {prompt_msg.lower()}?' if ui.ask(msg): # Replace part list with whole disk obj object_list = [disk.path] @@ -2387,11 +2377,11 @@ def select_disk_parts(prompt, disk): return object_list -def select_path(prompt): +def select_path(prompt_msg): """Select path, returns pathlib.Path.""" invalid = False menu = ui.Menu( - title=ansi.color_string(f'ddrescue TUI: {prompt} Path Selection', 'GREEN'), + title=ansi.color_string(f'ddrescue TUI: {prompt_msg} Path Selection', 'GREEN'), ) menu.separator = ' ' menu.add_action('Quit') @@ -2433,7 +2423,7 @@ def set_mode(docopt_args): # Ask user if necessary if not mode: - answer = ui.choice(['C', 'I'], 'Are we cloning or imaging?') + answer = ui.choice('Are we cloning or imaging?', ['C', 'I']) if answer == 'C': mode = 'Clone' else: diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py index 3b2b9d93..2cdb1df9 100644 --- a/scripts/wk/os/win.py +++ b/scripts/wk/os/win.py @@ -212,7 +212,7 @@ def export_bitlocker_info(): ] # Get filename - file_name = ui.input_text(prompt='Enter filename', allow_empty_response=False) + file_name = ui.input_text(prompt_msg='Enter filename') file_path = pathlib.Path(f'../../Bitlocker_{file_name}.txt').resolve() # Save info diff --git a/scripts/wk/ui/cli.py b/scripts/wk/ui/cli.py index 298fecaf..83cc2ed0 100644 --- a/scripts/wk/ui/cli.py +++ b/scripts/wk/ui/cli.py @@ -10,6 +10,8 @@ import sys import traceback from collections import OrderedDict +from prompt_toolkit import prompt +from prompt_toolkit.validation import Validator, ValidationError try: from functools import cache @@ -32,6 +34,74 @@ PLATFORM = platform.system() # Classes +class InputChoiceValidator(Validator): + """Validate that input is one of the provided choices.""" + def __init__(self, choices, allow_empty=False): + self.allow_empty = allow_empty + self.choices = [str(c).upper() for c in choices] + super().__init__() + + def validate(self, document): + text = document.text + if not (text or self.allow_empty): + raise ValidationError( + message='This input is required!', + cursor_position=len(text), + ) + if text and text.upper() not in self.choices: + raise ValidationError( + message='Invalid selection', + cursor_position=len(text), + ) + +class InputNotEmptyValidator(Validator): + """Validate that input is not empty.""" + def validate(self, document): + text = document.text + if not text: + raise ValidationError( + message='This input is required!', + cursor_position=len(text), + ) + +class InputTicketIDValidator(Validator): + """Validate that input resembles a ticket ID.""" + def __init__(self, allow_empty=False): + self.allow_empty = allow_empty + super().__init__() + + def validate(self, document): + text = document.text + if not (text or self.allow_empty): + raise ValidationError( + message='This input is required!', + cursor_position=len(text), + ) + if text and not re.match(r'^\d', text): + raise ValidationError( + message='Ticket ID should start with a number!', + cursor_position=len(text), + ) + +class InputYesNoValidator(Validator): + """Validate that input is a yes or no.""" + def __init__(self, allow_empty=False): + self.allow_empty = allow_empty + super().__init__() + + def validate(self, document): + text = document.text + if not (text or self.allow_empty): + raise ValidationError( + message='This input is required!', + cursor_position=len(text), + ) + if text and not re.match(r'^(y(es|up|)|n(o|ope|))$', text, re.IGNORECASE): + raise ValidationError( + message='Please answer "yes" or "no"', + cursor_position=len(text), + ) + class Menu(): """Object for tracking menu specific data and methods. @@ -250,7 +320,7 @@ class Menu(): # Otherwise deselect details['Selected'] = status and option in targets - def _user_select(self, prompt): + def _user_select(self, prompt_msg): """Show menu and select an entry, returns str.""" menu_text = self._generate_menu_text() valid_answers = self._get_valid_answers() @@ -260,7 +330,7 @@ class Menu(): clear_screen() print(menu_text) sleep(0.01) - answer = input_text(prompt).strip() + answer = input_text(prompt_msg).strip() if answer.upper() in valid_answers: break @@ -297,14 +367,14 @@ class Menu(): details['Selected'] = details.get('Selected', False) self.toggles[name] = details - def advanced_select(self, prompt='Please make a selection: '): + def advanced_select(self, prompt_msg='Please make a selection: '): """Display menu and make multiple selections, returns tuple. NOTE: Menu is displayed until an action entry is selected. """ while True: self._update(single_selection=False) - user_selection = self._user_select(prompt) + user_selection = self._user_select(prompt_msg) selected_entry = self._resolve_selection(user_selection) if user_selection.isnumeric(): # Update selection(s) @@ -316,19 +386,19 @@ class Menu(): # Done return selected_entry - def settings_select(self, prompt='Please make a selection: '): + def settings_select(self, prompt_msg='Please make a selection: '): """Display menu and make multiple selections, returns tuple. NOTE: Menu is displayed until an action entry is selected. """ choice_kwargs = { + 'prompt_msg': 'Toggle or change value?', 'choices': ['T', 'C'], - 'prompt': 'Toggle or change value?', } while True: self._update(single_selection=True, settings_mode=True) - user_selection = self._user_select(prompt) + user_selection = self._user_select(prompt_msg) selected_entry = self._resolve_selection(user_selection) if user_selection.isnumeric(): if 'Value' in selected_entry[-1] and choice(**choice_kwargs) == 'C': @@ -344,18 +414,17 @@ class Menu(): # Done return selected_entry - def simple_select(self, prompt='Please make a selection: ', update=True): + def simple_select(self, prompt_msg='Please make a selection: ', update=True): """Display menu and make a single selection, returns tuple.""" if update: self._update() - user_selection = self._user_select(prompt) + user_selection = self._user_select(prompt_msg) return self._resolve_selection(user_selection) def update(self): """Update menu with default settings.""" self._update() - class TryAndPrint(): """Object used to standardize running functions and returning the result. @@ -563,30 +632,28 @@ class TryAndPrint(): # Functions -def abort(prompt='Aborted.', show_prompt=True, return_code=1): +def abort(prompt_msg='Aborted.', show_prompt_msg=True, return_code=1): """Abort script.""" - print_warning(prompt) - if show_prompt: + print_warning(prompt_msg) + if show_prompt_msg: sleep(0.5) - pause(prompt='Press Enter to exit... ') + pause(prompt_msg='Press Enter to exit... ') sys.exit(return_code) -def ask(prompt='Kotaero!'): +def ask(prompt_msg): """Prompt the user with a Y/N question, returns bool.""" - answer = None - prompt = f'{prompt} [Y/N]: ' + validator = InputYesNoValidator() - # Loop until acceptable answer is given - while answer is None: - tmp = input_text(prompt) - if re.search(r'^y(es|up|)$', tmp, re.IGNORECASE): - answer = True - elif re.search(r'^n(o|ope|)$', tmp, re.IGNORECASE): - answer = False + # Show prompt + response = input_text(f'{prompt_msg} [Y/N]: ', validator=validator) + if response.upper().startswith('Y'): + answer = True + elif response.upper().startswith('N'): + answer = False # Done - LOG.info('%s%s', prompt, 'Yes' if answer else 'No') + LOG.info('%s%s', prompt_msg, 'Yes' if answer else 'No') return answer @@ -599,27 +666,32 @@ def beep(repeat=1): repeat -= 1 -def choice(choices, prompt='答えろ！'): +def choice(prompt_msg, choices): """Choose an option from a provided list, returns str. Choices provided will be converted to uppercase and returned as such. Similar to the commands choice (Windows) and select (Linux). """ - LOG.debug('choices: %s, prompt: %s', choices, prompt) - answer = None + LOG.debug('prompt_msg: %s, choices: %s', prompt_msg, choices) choices = [str(c).upper()[:1] for c in choices] - prompt = f'{prompt} [{"/".join(choices)}]' - regex = f'^({"|".join(choices)})$' + prompt_msg = f'{prompt_msg} [{"/".join(choices)}]' - # Loop until acceptable answer is given - while answer is None: - tmp = input_text(prompt=prompt) - if re.search(regex, tmp, re.IGNORECASE): - answer = tmp.upper() + # Show prompt + response = input_text(prompt_msg, validator=InputChoiceValidator(choices)) # Done - LOG.info('%s %s', prompt, answer) - return answer + LOG.info('%s %s', prompt_msg, response) + return response.upper() + + +def fix_prompt(message): + """Fix prompt, returns str.""" + if not message: + message = 'Input text: ' + message = str(message) + if message[-1:] != ' ': + message += ' ' + return message @cache @@ -659,29 +731,39 @@ def get_exception(name): return obj -def input_text(prompt='Enter text', allow_empty_response=True): - """Get text from user, returns string.""" - prompt = str(prompt) - response = None - if prompt[-1:] != ' ': - prompt += ' ' - print(prompt, end='', flush=True) +def get_ticket_id(): + """Get ticket ID, returns str.""" + prompt_msg = 'Please enter ticket ID:' + validator = InputTicketIDValidator() - while response is None: + # Show prompt + ticket_id = input_text(prompt_msg, validator=validator) + + # Done + return ticket_id + + +def input_text( + prompt_msg='Enter text: ', allow_empty=False, validator=None, + ) -> str: + """Get input from user, returns str.""" + prompt_msg = fix_prompt(prompt_msg) + + # Accept empty responses? + if not (allow_empty or validator): + validator = InputNotEmptyValidator() + + # Show prompt + result = None + while result is None: try: - response = input() - LOG.debug('%s%s', prompt, response) - except EOFError: - # Ignore and try again - LOG.warning('Exception occured', exc_info=True) - print('', flush=True) - if not allow_empty_response: - if response is None or not response.strip(): - # The None check here is used to avoid a TypeError if response is None - print(f'\r{prompt}', end='', flush=True) - response = None + result = prompt(prompt_msg, validator=validator) + except KeyboardInterrupt: + # Ignore CTRL+c + pass - return response + # Done + return result def major_exception(): @@ -698,9 +780,9 @@ def major_exception(): raise SystemExit(1) -def pause(prompt='Press Enter to continue... '): +def pause(prompt_msg='Press Enter to continue... '): """Simple pause implementation.""" - input_text(prompt) + input_text(prompt_msg, allow_empty=True) def print_colored(strings, colors, log=False, sep=' ', **kwargs): diff --git a/setup/windows/build.ps1 b/setup/windows/build.ps1 index b7ed43a7..a7278265 100644 --- a/setup/windows/build.ps1 +++ b/setup/windows/build.ps1 @@ -107,9 +107,9 @@ if ($MyInvocation.InvocationName -ne ".") { $Url = FindDynamicUrl $DownloadPage $RegEx DownloadFile -Path $Temp -Name "psutil64.whl" -Url $Url - # Python: pytz, requests, & dependancies + # Python: prompt_toolkit, pytz, requests, & dependancies $RegEx = "href=.*.py3-none-any.whl" - foreach ($Module in @("chardet", "certifi", "idna", "pytz", "urllib3", "requests")) { + foreach ($Module in @("chardet", "certifi", "idna", "prompt_toolkit", "Pygments", "pytz", "requests", "urllib3", "wcwidth")) { $DownloadPage = "https://pypi.org/project/$Module/" $Name = "$Module.whl" $Url = FindDynamicUrl -SourcePage $DownloadPage -RegEx $RegEx From f9bcd534d441000770176b70e37e1fec4e3240ac Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sat, 8 Apr 2023 19:20:22 -0700 Subject: [PATCH 08/16] Move ansi code into wk/ui --- scripts/wk/__init__.py | 1 - scripts/wk/clone/ddrescue.py | 4 ++-- scripts/wk/graph.py | 2 +- scripts/wk/hw/benchmark.py | 3 ++- scripts/wk/hw/cpu.py | 4 ++-- scripts/wk/hw/diags.py | 4 ++-- scripts/wk/hw/disk.py | 2 +- scripts/wk/hw/sensors.py | 2 +- scripts/wk/hw/smart.py | 2 +- scripts/wk/hw/surface_scan.py | 2 +- scripts/wk/hw/system.py | 2 +- scripts/wk/os/linux.py | 2 +- scripts/wk/os/win.py | 2 +- scripts/wk/repairs/win.py | 4 ++-- scripts/wk/setup/win.py | 4 ++-- scripts/wk/ui/__init__.py | 1 + scripts/wk/{ => ui}/ansi.py | 0 scripts/wk/ui/cli.py | 2 +- 18 files changed, 22 insertions(+), 21 deletions(-) rename scripts/wk/{ => ui}/ansi.py (100%) diff --git a/scripts/wk/__init__.py b/scripts/wk/__init__.py index 2a77ed94..afb3e8d5 100644 --- a/scripts/wk/__init__.py +++ b/scripts/wk/__init__.py @@ -3,7 +3,6 @@ from sys import stderr, version_info -from . import ansi from . import cfg from . import clone from . import debug diff --git a/scripts/wk/clone/ddrescue.py b/scripts/wk/clone/ddrescue.py index 2c8bd568..1bc96cff 100644 --- a/scripts/wk/clone/ddrescue.py +++ b/scripts/wk/clone/ddrescue.py @@ -20,7 +20,7 @@ from docopt import docopt import psutil import pytz -from wk import ansi, cfg, debug, exe, io, log, net, std +from wk import cfg, debug, exe, io, log, net, std from wk.cfg.ddrescue import ( DDRESCUE_MAP_TEMPLATE, DDRESCUE_SETTINGS, @@ -34,7 +34,7 @@ from wk.hw.smart import ( update_smart_details, ) from wk.ui import cli as ui -from wk.ui import tmux +from wk.ui import ansi, tmux # STATIC VARIABLES diff --git a/scripts/wk/graph.py b/scripts/wk/graph.py index 70e26cfb..95324956 100644 --- a/scripts/wk/graph.py +++ b/scripts/wk/graph.py @@ -3,7 +3,7 @@ import logging -from wk import ansi +from wk.ui import ansi # STATIC VARIABLES diff --git a/scripts/wk/hw/benchmark.py b/scripts/wk/hw/benchmark.py index 5a33ad88..71a0a465 100644 --- a/scripts/wk/hw/benchmark.py +++ b/scripts/wk/hw/benchmark.py @@ -5,7 +5,7 @@ import logging from subprocess import PIPE, STDOUT -from wk import ansi, graph +from wk import graph from wk.cfg.hw import ( IO_ALT_TEST_SIZE_FACTOR, IO_BLOCK_SIZE, @@ -22,6 +22,7 @@ from wk.cfg.hw import ( ) from wk.exe import run_program from wk.std import PLATFORM +from wk.ui import ansi # STATIC VARIABLES diff --git a/scripts/wk/hw/cpu.py b/scripts/wk/hw/cpu.py index deb8f808..18a4bd83 100644 --- a/scripts/wk/hw/cpu.py +++ b/scripts/wk/hw/cpu.py @@ -7,11 +7,11 @@ import subprocess from typing import TextIO -from wk import ansi, exe +from wk import exe from wk.cfg.hw import CPU_FAILURE_TEMP from wk.os.mac import set_fans as macos_set_fans from wk.std import PLATFORM -from wk.ui import tmux +from wk.ui import ansi, tmux # STATIC VARIABLES diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py index 0ad9206c..dc214205 100644 --- a/scripts/wk/hw/diags.py +++ b/scripts/wk/hw/diags.py @@ -10,7 +10,7 @@ import time from docopt import docopt -from wk import ansi, cfg, debug, exe, log, std +from wk import cfg, debug, exe, log, std from wk.cfg.hw import STATUS_COLORS from wk.hw import benchmark as hw_benchmark from wk.hw import cpu as hw_cpu @@ -26,7 +26,7 @@ from wk.hw.screensavers import screensaver from wk.hw.test import Test, TestGroup from wk.ui import cli as ui -from wk.ui import tmux +from wk.ui import ansi, tmux # STATIC VARIABLES diff --git a/scripts/wk/hw/disk.py b/scripts/wk/hw/disk.py index bae3e305..31704c8f 100644 --- a/scripts/wk/hw/disk.py +++ b/scripts/wk/hw/disk.py @@ -11,7 +11,6 @@ import re from dataclasses import dataclass, field from typing import Any, Union -from wk import ansi from wk.cfg.main import KIT_NAME_SHORT from wk.cfg.python import DATACLASS_DECORATOR_KWARGS from wk.exe import get_json_from_command, run_program @@ -21,6 +20,7 @@ from wk.hw.smart import ( get_known_disk_attributes, ) from wk.std import PLATFORM +from wk.ui import ansi # STATIC VARIABLES diff --git a/scripts/wk/hw/sensors.py b/scripts/wk/hw/sensors.py index f05a628d..815a6de5 100644 --- a/scripts/wk/hw/sensors.py +++ b/scripts/wk/hw/sensors.py @@ -9,11 +9,11 @@ import re from subprocess import CalledProcessError from typing import Any -from wk import ansi from wk.cfg.hw import CPU_CRITICAL_TEMP, SMC_IDS, TEMP_COLORS from wk.exe import run_program, start_thread from wk.io import non_clobber_path from wk.std import PLATFORM, sleep +from wk.ui import ansi # STATIC VARIABLES diff --git a/scripts/wk/hw/smart.py b/scripts/wk/hw/smart.py index 88c016ab..6a099701 100644 --- a/scripts/wk/hw/smart.py +++ b/scripts/wk/hw/smart.py @@ -7,7 +7,6 @@ import re from typing import Any -from wk import ansi from wk.cfg.hw import ( ATTRIBUTE_COLORS, KEY_NVME, @@ -20,6 +19,7 @@ from wk.cfg.hw import ( ) from wk.exe import get_json_from_command, run_program from wk.std import bytes_to_string, sleep +from wk.ui import ansi # STATIC VARIABLES diff --git a/scripts/wk/hw/surface_scan.py b/scripts/wk/hw/surface_scan.py index 7f4ed243..135c9951 100644 --- a/scripts/wk/hw/surface_scan.py +++ b/scripts/wk/hw/surface_scan.py @@ -5,7 +5,6 @@ import logging from subprocess import STDOUT -from wk import ansi from wk.cfg.hw import ( BADBLOCKS_EXTRA_LARGE_DISK, BADBLOCKS_LARGE_DISK, @@ -16,6 +15,7 @@ from wk.cfg.hw import ( ) from wk.exe import run_program from wk.std import PLATFORM, bytes_to_string +from wk.ui import ansi # STATIC VARIABLES diff --git a/scripts/wk/hw/system.py b/scripts/wk/hw/system.py index 949adb38..4309c3df 100644 --- a/scripts/wk/hw/system.py +++ b/scripts/wk/hw/system.py @@ -8,12 +8,12 @@ import re from dataclasses import dataclass, field from typing import Any -from wk import ansi from wk.cfg.hw import KNOWN_RAM_VENDOR_IDS from wk.cfg.python import DATACLASS_DECORATOR_KWARGS from wk.exe import get_json_from_command, run_program from wk.hw.test import Test from wk.std import PLATFORM, bytes_to_string, string_to_bytes +from wk.ui import ansi # STATIC VARIABLES diff --git a/scripts/wk/os/linux.py b/scripts/wk/os/linux.py index a4c82002..ee17792a 100644 --- a/scripts/wk/os/linux.py +++ b/scripts/wk/os/linux.py @@ -7,11 +7,11 @@ import pathlib import re import subprocess -from wk import ansi from wk.cfg.hw import VOLUME_FAILURE_THRESHOLD, VOLUME_WARNING_THRESHOLD from wk.exe import get_json_from_command, popen_program, run_program from wk.log import format_log_path from wk.std import bytes_to_string +from wk.ui import ansi # STATIC VARIABLES diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py index 2cdb1df9..4d6d7248 100644 --- a/scripts/wk/os/win.py +++ b/scripts/wk/os/win.py @@ -17,7 +17,6 @@ except ImportError as err: if platform.system() == 'Windows': raise err -from wk import ansi from wk.borrowed import acpi from wk.cfg.main import KIT_NAME_FULL from wk.cfg.windows_builds import ( @@ -34,6 +33,7 @@ from wk.std import ( sleep, ) from wk.ui import cli as ui +from wk.ui import ansi # STATIC VARIABLES diff --git a/scripts/wk/repairs/win.py b/scripts/wk/repairs/win.py index 5ae419ed..51c418d4 100644 --- a/scripts/wk/repairs/win.py +++ b/scripts/wk/repairs/win.py @@ -11,7 +11,6 @@ import time from subprocess import CalledProcessError, DEVNULL from xml.dom.minidom import parse as xml_parse -from wk import ansi from wk.cfg.main import KIT_NAME_FULL, KIT_NAME_SHORT, WINDOWS_TIME_ZONE from wk.cfg.repairs import ( AUTO_REPAIR_DELAY_IN_SECONDS, @@ -61,7 +60,8 @@ from wk.std import ( GenericWarning, sleep, ) -from wk.ui import cli as ui +from wk.ui import cli as ui +from wk.ui import ansi # STATIC VARIABLES diff --git a/scripts/wk/setup/win.py b/scripts/wk/setup/win.py index 96157af3..97f93b5b 100644 --- a/scripts/wk/setup/win.py +++ b/scripts/wk/setup/win.py @@ -8,7 +8,6 @@ import os import re import sys -from wk import ansi from wk.cfg.main import KIT_NAME_FULL from wk.cfg.setup import ( BROWSER_PATHS, @@ -63,7 +62,8 @@ from wk.std import ( GenericWarning, sleep, ) -from wk.ui import cli as ui +from wk.ui import cli as ui +from wk.ui import ansi # STATIC VARIABLES diff --git a/scripts/wk/ui/__init__.py b/scripts/wk/ui/__init__.py index 209a7699..b93cb848 100644 --- a/scripts/wk/ui/__init__.py +++ b/scripts/wk/ui/__init__.py @@ -1,4 +1,5 @@ """WizardKit: ui module init""" +from . import ansi from . import cli from . import tmux diff --git a/scripts/wk/ansi.py b/scripts/wk/ui/ansi.py similarity index 100% rename from scripts/wk/ansi.py rename to scripts/wk/ui/ansi.py diff --git a/scripts/wk/ui/cli.py b/scripts/wk/ui/cli.py index 83cc2ed0..3fb7c173 100644 --- a/scripts/wk/ui/cli.py +++ b/scripts/wk/ui/cli.py @@ -19,7 +19,6 @@ except ImportError: # Assuming Python is < 3.9 from functools import lru_cache as cache -from wk.ansi import clear_screen, color_string, strip_colors from wk.cfg.main import ( ENABLED_UPLOAD_DATA, INDENT, @@ -27,6 +26,7 @@ from wk.cfg.main import ( WIDTH, ) from wk.std import (sleep, GenericWarning) +from wk.ui.ansi import clear_screen, color_string, strip_colors # STATIC VARIABLES LOG = logging.getLogger(__name__) From ba3bf480f7964e033278e981ace84b1037b8d8d6 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sun, 9 Apr 2023 15:59:34 -0700 Subject: [PATCH 09/16] BREAKING Add wk/ui/tui.py --- scripts/wk/cfg/hw.py | 14 -- scripts/wk/hw/diags.py | 153 ++++++++----------- scripts/wk/ui/__init__.py | 1 + scripts/wk/ui/tmux.py | 124 ++++++++++------ scripts/wk/ui/tui.py | 300 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 448 insertions(+), 144 deletions(-) create mode 100644 scripts/wk/ui/tui.py diff --git a/scripts/wk/cfg/hw.py b/scripts/wk/cfg/hw.py index 4ebadb84..2ebfd724 100644 --- a/scripts/wk/cfg/hw.py +++ b/scripts/wk/cfg/hw.py @@ -3,8 +3,6 @@ import re -from collections import OrderedDict - # STATIC VARIABLES ATTRIBUTE_COLORS = ( @@ -161,18 +159,6 @@ THRESH_HDD_AVG_LOW = 65 * 1024**2 THRESH_SSD_MIN = 90 * 1024**2 THRESH_SSD_AVG_HIGH = 135 * 1024**2 THRESH_SSD_AVG_LOW = 100 * 1024**2 -TMUX_SIDE_WIDTH = 20 -TMUX_LAYOUT = OrderedDict({ - 'Top': {'height': 2, 'Check': True}, - 'Started': {'width': TMUX_SIDE_WIDTH, 'Check': True}, - 'Progress': {'width': TMUX_SIDE_WIDTH, 'Check': True}, - # Testing panes - 'Temps': {'height': 1000, 'Check': False}, - 'Prime95': {'height': 11, 'Check': False}, - 'SMART': {'height': 4, 'Check': True}, - 'badblocks': {'height': 5, 'Check': True}, - 'I/O Benchmark': {'height': 1000, 'Check': False}, - }) # VOLUME THRESHOLDS in percent VOLUME_WARNING_THRESHOLD = 70 VOLUME_FAILURE_THRESHOLD = 85 diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py index dc214205..60e42abe 100644 --- a/scripts/wk/hw/diags.py +++ b/scripts/wk/hw/diags.py @@ -25,8 +25,8 @@ from wk.hw.network import network_test from wk.hw.screensavers import screensaver from wk.hw.test import Test, TestGroup -from wk.ui import cli as ui -from wk.ui import ansi, tmux +from wk.ui import tui as ui +from wk.ui import ansi, cli, tmux # STATIC VARIABLES @@ -84,18 +84,15 @@ class State(): """Object for tracking hardware diagnostic data.""" def __init__(self, test_mode=False): self.disks = [] - self.layout = cfg.hw.TMUX_LAYOUT.copy() self.log_dir = None self.panes = {} + self.progress_file = None self.system = None self.test_groups = [] - self.top_text = ansi.color_string('Hardware Diagnostics', 'GREEN') + self.title_text = ansi.color_string('Hardware Diagnostics', 'GREEN') if test_mode: - self.top_text += ansi.color_string(' (Test Mode)', 'YELLOW') - - # Init tmux and start a background process to maintain layout - self.init_tmux() - exe.start_thread(self.fix_tmux_layout_loop) + self.title_text += ansi.color_string(' (Test Mode)', 'YELLOW') + self.ui = ui.TUI(self.title_text) def abort_testing(self) -> None: """Set unfinished tests as aborted and cleanup tmux panes.""" @@ -125,23 +122,6 @@ class State(): dev.disable_disk_tests() break - 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) - except RuntimeError: - # Assuming self.panes changed while running - pass - - def fix_tmux_layout_loop(self) -> None: - """Fix tmux layout on a loop. - - NOTE: This should be called as a thread. - """ - while True: - self.fix_tmux_layout(forced=False) - std.sleep(1) - def init_diags(self, menu) -> None: """Initialize diagnostic pass.""" @@ -163,15 +143,13 @@ class State(): keep_history=False, timestamp=False, ) - ui.clear_screen() - ui.print_info('Initializing...') + cli.clear_screen() + cli.print_info('Initializing...') # Progress Pane - self.update_progress_pane() - tmux.respawn_pane( - pane_id=self.panes['Progress'], - watch_file=f'{self.log_dir}/progress.out', - ) + self.progress_file = pathlib.Path(f'{self.log_dir}/progress.out') + self.update_progress_file() + self.ui.set_progress_file(self.progress_file) # Add HW Objects self.system = hw_system.System() @@ -222,7 +200,7 @@ class State(): behind=True, lines=2, vertical=True, - text=f'{self.top_text}\nMain Menu', + text=f'{self.title_text}\nMain Menu', ) # Started @@ -299,16 +277,15 @@ class State(): ), ) - def update_progress_pane(self) -> None: - """Update progress pane.""" + def update_progress_file(self) -> None: + """Update progress file.""" report = [] - width = cfg.hw.TMUX_SIDE_WIDTH for group in self.test_groups: report.append(ansi.color_string(group.name, 'BLUE')) for test in group.test_objects: report.append(ansi.color_string( - [test.label, f'{test.status:>{width-len(test.label)}}'], + [test.label, f'{test.status:>{self.ui.side_width-len(test.label)}}'], [None, STATUS_COLORS.get(test.status, None)], sep='', )) @@ -317,19 +294,17 @@ class State(): report.append(' ') # Write to progress file - out_path = pathlib.Path(f'{self.log_dir}/progress.out') - with open(out_path, 'w', encoding='utf-8') as _f: - _f.write('\n'.join(report)) + self.progress_file.write_text('\n'.join(report), encoding='utf-8') 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}') + tmux.respawn_pane(self.panes['Top'], text=f'{self.title_text}\n{text}') # Functions -def build_menu(cli_mode=False, quick_mode=False) -> ui.Menu: +def build_menu(cli_mode=False, quick_mode=False) -> cli.Menu: """Build main menu, returns wk.ui.cli.Menu.""" - menu = ui.Menu(title=None) + menu = cli.Menu(title=None) # Add actions, options, etc for action in MENU_ACTIONS: @@ -407,7 +382,7 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: ) # Create monitor and worker panes - state.update_progress_pane() + state.update_progress_file() state.panes['Prime95'] = tmux.split_window( lines=10, vertical=True, watch_file=prime_log, watch_cmd='tail') if PLATFORM == 'Darwin': @@ -421,11 +396,11 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: state.layout['Current'] = {'height': 3, 'Check': True} # Get idle temps - ui.print_standard('Saving idle temps...') + cli.print_standard('Saving idle temps...') sensors.save_average_temps(temp_label='Idle', seconds=5) # Stress CPU - ui.print_info('Running stress test') + cli.print_info('Running stress test') hw_cpu.set_apple_fan_speed('max') proc_mprime = hw_cpu.start_mprime(state.log_dir, prime_log) @@ -443,13 +418,13 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: if sensors.cpu_reached_critical_temp() or aborted: test_cooling_obj.set_status('Aborted') test_mprime_obj.set_status('Aborted') - state.update_progress_pane() + state.update_progress_file() # Get cooldown temp - ui.clear_screen() - ui.print_standard('Letting CPU cooldown...') + cli.clear_screen() + cli.print_standard('Letting CPU cooldown...') std.sleep(5) - ui.print_standard('Saving cooldown temps...') + cli.print_standard('Saving cooldown temps...') sensors.save_average_temps(temp_label='Cooldown', seconds=5) # Check Prime95 results @@ -464,10 +439,10 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: ) if run_sysbench: LOG.info('CPU Test (Sysbench)') - ui.print_standard('Letting CPU cooldown more...') + cli.print_standard('Letting CPU cooldown more...') std.sleep(30) - ui.clear_screen() - ui.print_info('Running alternate stress test') + cli.clear_screen() + cli.print_info('Running alternate stress test') print('') proc_sysbench, filehandle_sysbench = hw_cpu.start_sysbench( sensors, @@ -490,14 +465,14 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: if sensors.cpu_reached_critical_temp() or aborted: test_cooling_obj.set_status('Aborted') test_mprime_obj.set_status('Aborted') - state.update_progress_pane() + state.update_progress_file() # Check Cooling results test_cooling_obj.report.append(ansi.color_string('Temps', 'BLUE')) hw_cpu.check_cooling_results(test_cooling_obj, sensors, run_sysbench) # Cleanup - state.update_progress_pane() + state.update_progress_file() sensors.stop_background_monitor() state.panes.pop('Current', None) tmux.kill_pane(state.panes.pop('Prime95', None)) @@ -519,7 +494,7 @@ def disk_attribute_check(state, test_objects, test_mode=False) -> None: continue # Done - state.update_progress_pane() + state.update_progress_file() def disk_io_benchmark( @@ -548,8 +523,8 @@ def disk_io_benchmark( continue # Start benchmark - ui.clear_screen() - ui.print_report(test.dev.generate_report()) + cli.clear_screen() + cli.print_report(test.dev.generate_report()) test.set_status('Working') test_log = f'{state.log_dir}/{test.dev.path.name}_benchmark.out' tmux.respawn_pane( @@ -557,7 +532,7 @@ def disk_io_benchmark( watch_cmd='tail', watch_file=test_log, ) - state.update_progress_pane() + state.update_progress_file() try: hw_benchmark.run_io_test(test, test_log, test_mode=test_mode) except KeyboardInterrupt: @@ -575,10 +550,10 @@ def disk_io_benchmark( break # Update progress after each test - state.update_progress_pane() + state.update_progress_file() # Cleanup - state.update_progress_pane() + state.update_progress_file() tmux.kill_pane(state.panes.pop('I/O Benchmark', None)) # Done @@ -597,7 +572,7 @@ def disk_self_test(state, test_objects, test_mode=False) -> None: state.update_top_pane( f'Disk self-test{"s" if len(test_objects) > 1 else ""}', ) - ui.print_info(f'Starting self-test{"s" if len(test_objects) > 1 else ""}') + cli.print_info(f'Starting self-test{"s" if len(test_objects) > 1 else ""}') show_failed_attributes(state) for test in reversed(test_objects): if test.disabled: @@ -616,7 +591,7 @@ def disk_self_test(state, test_objects, test_mode=False) -> None: ) # Wait for all tests to complete - state.update_progress_pane() + state.update_progress_file() try: while True: if any(t.is_alive() for t in threads): @@ -631,7 +606,7 @@ def disk_self_test(state, test_objects, test_mode=False) -> None: hw_smart.build_self_test_report(test, aborted=True) # Cleanup - state.update_progress_pane() + state.update_progress_file() for pane in state.panes['SMART']: tmux.kill_pane(pane) state.panes.pop('SMART', None) @@ -694,7 +669,7 @@ def disk_surface_scan(state, test_objects, test_mode=False) -> None: state.update_top_pane( f'Disk Surface Scan{"s" if len(test_objects) > 1 else ""}', ) - ui.print_info( + cli.print_info( f'Starting disk surface scan{"s" if len(test_objects) > 1 else ""}', ) show_failed_attributes(state) @@ -723,7 +698,7 @@ def disk_surface_scan(state, test_objects, test_mode=False) -> None: try: while True: if any(t.is_alive() for t in threads): - state.update_progress_pane() + state.update_progress_file() std.sleep(5) else: break @@ -737,7 +712,7 @@ def disk_surface_scan(state, test_objects, test_mode=False) -> None: test.report.append(ansi.color_string(' Aborted', 'YELLOW')) # Cleanup - state.update_progress_pane() + state.update_progress_file() for pane in state.panes['badblocks']: tmux.kill_pane(pane) state.panes.pop('badblocks', None) @@ -788,9 +763,9 @@ def main() -> None: try: action() except KeyboardInterrupt: - ui.print_warning('Aborted.') - ui.print_standard('') - ui.pause('Press Enter to return to main menu...') + cli.print_warning('Aborted.') + cli.print_standard('') + cli.pause('Press Enter to return to main menu...') if 'Clock Sync' in selection: state.update_clock() @@ -855,8 +830,8 @@ def run_diags(state, menu, quick_mode=False, test_mode=False) -> None: # Just return if no tests were selected if not state.test_groups: - ui.print_warning('No tests selected?') - ui.pause() + cli.print_warning('No tests selected?') + cli.pause() return # Run tests @@ -867,13 +842,13 @@ def run_diags(state, menu, quick_mode=False, test_mode=False) -> None: args = [group.test_objects] if group.name == 'Disk I/O Benchmark': args.append(menu.toggles['Skip USB Benchmarks']['Selected']) - ui.clear_screen() + cli.clear_screen() try: function(state, *args, test_mode=test_mode) except (KeyboardInterrupt, std.GenericAbort): aborted = True state.abort_testing() - state.update_progress_pane() + state.update_progress_file() break else: # Run safety checks after disk tests @@ -894,25 +869,25 @@ def run_diags(state, menu, quick_mode=False, test_mode=False) -> None: state.save_debug_reports() atexit.unregister(state.save_debug_reports) if quick_mode: - ui.pause('Press Enter to exit...') + cli.pause('Press Enter to exit...') else: - ui.pause('Press Enter to return to main menu...') + cli.pause('Press Enter to return to main menu...') def show_failed_attributes(state) -> None: """Show failed attributes for all disks.""" for dev in state.disks: - ui.print_colored([dev.name, dev.description], ['CYAN', None]) - ui.print_report( + cli.print_colored([dev.name, dev.description], ['CYAN', None]) + cli.print_report( hw_smart.generate_attribute_report(dev, only_failed=True), ) - ui.print_standard('') + cli.print_standard('') def show_results(state) -> None: """Show test results by device.""" std.sleep(0.5) - ui.clear_screen() + cli.clear_screen() state.update_top_pane('Results') # CPU Tests @@ -920,22 +895,22 @@ def show_results(state) -> None: group.name for group in state.test_groups if 'CPU' in group.name ] if cpu_tests_enabled: - ui.print_success('CPU:') - ui.print_report(state.system.generate_report()) - ui.print_standard(' ') + cli.print_success('CPU:') + cli.print_report(state.system.generate_report()) + cli.print_standard(' ') # Disk Tests disk_tests_enabled = [ group.name for group in state.test_groups if 'Disk' in group.name ] if disk_tests_enabled: - ui.print_success(f'Disk{"s" if len(state.disks) > 1 else ""}:') + cli.print_success(f'Disk{"s" if len(state.disks) > 1 else ""}:') for disk in state.disks: - ui.print_report(disk.generate_report()) - ui.print_standard(' ') + cli.print_report(disk.generate_report()) + cli.print_standard(' ') if not state.disks: - ui.print_warning('No devices') - ui.print_standard(' ') + cli.print_warning('No devices') + cli.print_standard(' ') def sync_clock() -> None: diff --git a/scripts/wk/ui/__init__.py b/scripts/wk/ui/__init__.py index b93cb848..75a335fc 100644 --- a/scripts/wk/ui/__init__.py +++ b/scripts/wk/ui/__init__.py @@ -3,3 +3,4 @@ from . import ansi from . import cli from . import tmux +from . import tui diff --git a/scripts/wk/ui/tmux.py b/scripts/wk/ui/tmux.py index 49fb3241..d473f81f 100644 --- a/scripts/wk/ui/tmux.py +++ b/scripts/wk/ui/tmux.py @@ -34,30 +34,76 @@ def clear_pane(pane_id=None): run_program(cmd, check=False) -def fix_layout(panes, layout, forced=False): +def fix_layout(layout, forced=False): """Fix pane sizes based on layout.""" - if not (forced or layout_needs_fixed(panes, layout)): + resize_kwargs = [] + + # Bail early + if not (forced or layout_needs_fixed(layout)): # Layout should be fine return - # Update panes - for name, data in layout.items(): - # Skip missing panes - if name not in panes: + # Remove closed panes + for data in layout.values(): + data['Panes'] = [pane for pane in data['Panes'] if poll_pane(pane)] + + # Update main panes + resize_pane(height=999) # Set active pane too large and then adjust down + for section, data in layout.items(): + if section == 'Workers': + # Skip for now continue - # Resize pane(s) - pane_list = panes[name] - if isinstance(pane_list, str): - pane_list = [pane_list] - for pane_id in pane_list: - if name == 'Current': - pane_id = None - try: - resize_pane(pane_id, **data) - except RuntimeError: - # Assuming pane was closed just before resizing - pass + if 'height' in data: + for pane_id in data['Panes']: + resize_kwargs.append({'pane_id': pane_id, 'height': data['height']}) + if 'width' in data: + for pane_id in data['Panes']: + resize_kwargs.append({'pane_id': pane_id, 'width': data['width']}) + for kwargs in resize_kwargs: + try: + resize_pane(**kwargs) + except RuntimeError: + # Assuming pane was closed just before resizing + pass + + # Update "group" panes widths + for group in ('Title', 'Info'): + num_panes = len(layout[group]['Panes']) + if num_panes <= 1: + continue + width = int( (get_pane_size()[0] - (1 - num_panes)) / num_panes ) + for pane_id in layout[group]['Panes']: + resize_pane(pane_id, width=width) + if group == 'Title': + # (re)fix Started pane + #TODO: REstore: resize_pane(layout['Started']['Panes'][0], width=TMUX_SIDE_WIDTH) + resize_pane(layout['Started']['Panes'][0], width=21) + + # Bail early + if not (layout['Workers']['Panes'] and layout['Workers']['height']): + return + + # Update worker heights + worker_height = layout['Workers']['height'] + workers = layout['Workers']['Panes'].copy() + num_workers = len(workers) + avail_height = sum(get_pane_size(pane)[1] for pane in workers) + avail_height += get_pane_size()[1] # Current pane + # Check if window is too small + if avail_height < (worker_height*num_workers) + 3: + # Just leave things as-is + return + # Resize current pane + resize_pane(height=avail_height-(worker_height*num_workers)) + # Resize bottom pane + resize_pane(workers.pop(0), height=worker_height) + # Resize the rest of the panes by adjusting the ones above them + while len(workers) > 1: + next_height = sum(get_pane_size(pane)[1] for pane in workers[:2]) + next_height -= worker_height + resize_pane(workers[1], height=next_height) + workers.pop(0) def get_pane_size(pane_id=None): @@ -96,34 +142,30 @@ def kill_pane(*pane_ids): run_program(cmd+[pane_id], check=False) -def layout_needs_fixed(panes, layout): +def layout_needs_fixed(layout): """Check if layout needs fixed, returns bool.""" needs_fixed = False # Check panes - for name, data in layout.items(): - # Skip unpredictably sized panes - if not data.get('Check', False): - continue + for data in layout.values(): + if 'height' in data: + needs_fixed = needs_fixed or any( + get_pane_size(pane)[1] != data['height'] for pane in data['Panes'] + ) + if 'width' in data: + needs_fixed = needs_fixed or any( + get_pane_size(pane)[0] != data['width'] for pane in data['Panes'] + ) - # Skip missing panes - if name not in panes: - continue - - # Check pane size(s) - pane_list = panes[name] - if isinstance(pane_list, str): - pane_list = [pane_list] - for pane_id in pane_list: - try: - width, height = get_pane_size(pane_id) - except ValueError: - # Pane may have disappeared during this loop - continue - if data.get('width', False) and data['width'] != width: - needs_fixed = True - if data.get('height', False) and data['height'] != height: - needs_fixed = True + # TODO: Re-enable? + ## Group panes + #for group in ('Title', 'Info'): + # num_panes = len(layout[group]['Panes']) + # if num_panes <= 1: + # continue + # width = int( (get_pane_size()[0] - (1 - num_panes)) / num_panes ) + # for pane in layout[group]['Panes']: + # needs_fixed = needs_fixed or abs(get_pane_size(pane)[0] - width) > 2 # Done return needs_fixed diff --git a/scripts/wk/ui/tui.py b/scripts/wk/ui/tui.py new file mode 100644 index 00000000..681f0a62 --- /dev/null +++ b/scripts/wk/ui/tui.py @@ -0,0 +1,300 @@ +"""WizardKit: TUI functions""" +# vim: sts=2 sw=2 ts=2 + +import logging +import time + +from copy import deepcopy + +from wk.exe import start_thread +from wk.std import sleep +from wk.ui import ansi, tmux + +# STATIC VARIABLES +LOG = logging.getLogger(__name__) +TMUX_SIDE_WIDTH = 21 +TMUX_TITLE_HEIGHT = 2 +TMUX_LAYOUT = { + 'Current': {'Panes': [None]}, + 'Title': {'Panes': [], 'height': TMUX_TITLE_HEIGHT}, + 'Started': {'Panes': [], 'width': TMUX_SIDE_WIDTH}, + 'Progress': {'Panes': [], 'width': TMUX_SIDE_WIDTH}, + 'Info': {'Panes': []}, + 'Workers': {'Panes': []}, + } + + +# Classes +class TUI(): + """Object for tracking TUI elements.""" + def __init__(self, title_text=None) -> None: + self.layout = deepcopy(TMUX_LAYOUT) + self.side_width = TMUX_SIDE_WIDTH + self.title_text = title_text if title_text else 'Title Text' + self.title_text_line2 = '' + self.title_colors = ['BLUE', None] + + # Init tmux and start a background process to maintain layout + self.init_tmux() + start_thread(self.fix_layout_loop) + + def add_info_pane(self, height, **tmux_args) -> None: + """Add info pane.""" + tmux_args.update({ + 'behind': True, + 'lines': height, + 'target_id': None, + 'vertical': True, + }) + if self.layout['Info']['Panes']: + tmux_args.update({ + 'behind': False, + 'percent': 50, + 'target_id': self.layout['Info']['Panes'][-1], + 'vertical': False, + }) + tmux_args.pop('lines') + + # Update layout + self.layout['Info']['height'] = height + + # Add pane + self.layout['Info']['Panes'].append(tmux.split_window(**tmux_args)) + + def add_title_pane(self, text) -> None: + """Add additional pane to title row.""" + tmux_args = { + 'behind': True, + 'lines': TMUX_TITLE_HEIGHT, + 'target_id': None, + 'text': text, + 'vertical': True, + } + if self.layout['Title']['Panes']: + tmux_args.update({ + 'behind': False, + 'percent': 50, + 'target_id': self.layout['Title']['Panes'][-1], + 'vertical': False, + }) + tmux_args.pop('lines') + + # Add pane + self.layout['Title']['Panes'].append(tmux.split_window(**tmux_args)) + + def add_worker_pane(self, height, **tmux_split_args) -> None: + """Add worker pane.""" + self.layout['Workers']['height'] = height + self.layout['Workers']['Panes'].append(tmux.split_window( + vertical=True, + lines=height, + **tmux_split_args, + )) + + def fix_layout(self, forced=True) -> None: + """Fix tmux layout based on self.layout.""" + try: + fix_layout(self.layout, forced=forced) + except RuntimeError: + # Assuming self.panes changed while running + pass + + def fix_layout_loop(self) -> None: + """Fix layout on a loop. + + NOTE: This should be called as a thread. + """ + while True: + self.fix_layout(forced=False) + sleep(1) + + def init_tmux(self) -> None: + """Initialize tmux layout.""" + tmux.kill_all_panes() + self.layout.clear() + self.layout.update(deepcopy(TMUX_LAYOUT)) + + # Title + self.layout['Title']['Panes'].append(tmux.split_window( + behind=True, + lines=2, + vertical=True, + text=ansi.color_string( + [self.title_text, self.title_text_line2], + self.title_colors, + sep = '\n', + ), + )) + + # Started + self.layout['Started']['Panes'].append(tmux.split_window( + lines=TMUX_SIDE_WIDTH, + target_id=self.layout['Title']['Panes'][0], + text=ansi.color_string( + ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], + ['BLUE', None], + sep='\n', + ), + )) + + # Progress + self.layout['Progress']['Panes'].append(tmux.split_window( + lines=TMUX_SIDE_WIDTH, + text=' ', + )) + + def remove_all_info_panes(self) -> None: + """Remove all info panes and update layout.""" + self.layout['Info'].pop('height', None) + panes = self.layout['Info']['Panes'].copy() + self.layout['Info']['Panes'].clear() + tmux.kill_pane(*panes) + + def remove_all_worker_panes(self) -> None: + """Remove all worker panes and update layout.""" + self.layout['Workers'].pop('height', None) + panes = self.layout['Workers']['Panes'].copy() + self.layout['Workers']['Panes'].clear() + tmux.kill_pane(*panes) + + def set_progress_file(self, progress_file) -> None: + """Set the file to use for the progresse pane.""" + tmux.respawn_pane( + pane_id=self.layout['Progress']['Panes'][0], + watch_file=progress_file, + ) + + def set_title(self, line1, line2=None, colors=None) -> None: + """Set title text.""" + self.title_text = line1 + self.title_text_line2 = line2 if line2 else '' + if colors: + self.title_colors = colors + + # Update pane (if present) + if self.layout['Title']['Panes']: + tmux.respawn_pane( + pane_id=self.layout['Title']['Panes'][0], + text=ansi.color_string( + [self.title_text, self.title_text_line2], + self.title_colors, + sep = '\n', + ), + ) + +# Functions +def fix_layout(layout, forced=False): + """Fix pane sizes based on layout.""" + resize_kwargs = [] + + # Bail early + if not (forced or layout_needs_fixed(layout)): + # Layout should be fine + return + + # Remove closed panes + for data in layout.values(): + data['Panes'] = [pane for pane in data['Panes'] if tmux.poll_pane(pane)] + + # Update main panes + tmux.resize_pane(height=999) # Set active pane too large and then adjust down + for section, data in layout.items(): + if section == 'Workers': + # Skip for now + continue + + if 'height' in data: + for pane_id in data['Panes']: + resize_kwargs.append({'pane_id': pane_id, 'height': data['height']}) + if 'width' in data: + for pane_id in data['Panes']: + resize_kwargs.append({'pane_id': pane_id, 'width': data['width']}) + for kwargs in resize_kwargs: + try: + tmux.resize_pane(**kwargs) + except RuntimeError: + # Assuming pane was closed just before resizing + pass + + # Update "group" panes widths + for group in ('Title', 'Info'): + num_panes = len(layout[group]['Panes']) + if num_panes <= 1: + continue + width = int( (tmux.get_pane_size()[0] - (1 - num_panes)) / num_panes ) + for pane_id in layout[group]['Panes']: + tmux.resize_pane(pane_id, width=width) + if group == 'Title': + # (re)fix Started pane + tmux.resize_pane(layout['Started']['Panes'][0], width=TMUX_SIDE_WIDTH) + + # Bail early + if not (layout['Workers']['Panes'] and layout['Workers']['height']): + return + + # Update worker heights + worker_height = layout['Workers']['height'] + workers = layout['Workers']['Panes'].copy() + num_workers = len(workers) + avail_height = sum(tmux.get_pane_size(pane)[1] for pane in workers) + avail_height += tmux.get_pane_size()[1] # Current pane + # Check if window is too small + if avail_height < (worker_height*num_workers) + 3: + # Just leave things as-is + return + # Resize current pane + tmux.resize_pane(height=avail_height-(worker_height*num_workers)) + # Resize bottom pane + tmux.resize_pane(workers.pop(0), height=worker_height) + # Resize the rest of the panes by adjusting the ones above them + while len(workers) > 1: + next_height = sum(tmux.get_pane_size(pane)[1] for pane in workers[:2]) + next_height -= worker_height + tmux.resize_pane(workers[1], height=next_height) + workers.pop(0) + +def layout_needs_fixed(layout): + """Check if layout needs fixed, returns bool.""" + needs_fixed = False + + # Check panes + for data in layout.values(): + if 'height' in data: + needs_fixed = needs_fixed or any( + tmux.get_pane_size(pane)[1] != data['height'] for pane in data['Panes'] + ) + if 'width' in data: + needs_fixed = needs_fixed or any( + tmux.get_pane_size(pane)[0] != data['width'] for pane in data['Panes'] + ) + + # TODO: Re-enable? + ## Group panes + #for group in ('Title', 'Info'): + # num_panes = len(layout[group]['Panes']) + # if num_panes <= 1: + # continue + # width = int( (tmux.get_pane_size()[0] - (1 - num_panes)) / num_panes ) + # for pane in layout[group]['Panes']: + # needs_fixed = needs_fixed or abs(tmux.get_pane_size(pane)[0] - width) > 2 + + # Done + return needs_fixed + +def hmm(): + """Hmm?""" + + +def test(): + ui = TUI() + ui.add_info_pane(height=10, text='Info One') + ui.add_info_pane(height=10, text='Info Two') + ui.add_info_pane(height=10, text='Info Three') + ui.add_worker_pane(height=3, text='Work One') + ui.add_worker_pane(height=3, text='Work Two') + ui.add_worker_pane(height=3, text='Work Three') + return ui + + +if __name__ == '__main__': + print("This file is not meant to be called directly.") From b834be9f00d920be70777ceaa76b3e5225ef96d8 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sun, 9 Apr 2023 16:32:13 -0700 Subject: [PATCH 10/16] Close all panes atexit --- scripts/wk/ui/tui.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/scripts/wk/ui/tui.py b/scripts/wk/ui/tui.py index 681f0a62..a6662d14 100644 --- a/scripts/wk/ui/tui.py +++ b/scripts/wk/ui/tui.py @@ -1,6 +1,7 @@ """WizardKit: TUI functions""" # vim: sts=2 sw=2 ts=2 +import atexit import logging import time @@ -38,6 +39,9 @@ class TUI(): self.init_tmux() start_thread(self.fix_layout_loop) + # Close all panes at exit + atexit.register(tmux.kill_all_panes) + def add_info_pane(self, height, **tmux_args) -> None: """Add info pane.""" tmux_args.update({ From 7aafcd7c01d0279334bf35e14d1a666102f9fc35 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sun, 9 Apr 2023 16:33:11 -0700 Subject: [PATCH 11/16] Update add_title_pane() to match set_title() --- scripts/wk/ui/tui.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/scripts/wk/ui/tui.py b/scripts/wk/ui/tui.py index a6662d14..f7834701 100644 --- a/scripts/wk/ui/tui.py +++ b/scripts/wk/ui/tui.py @@ -65,13 +65,18 @@ class TUI(): # Add pane self.layout['Info']['Panes'].append(tmux.split_window(**tmux_args)) - def add_title_pane(self, text) -> None: - """Add additional pane to title row.""" + def add_title_pane(self, line1, line2=None, colors=None) -> None: + """Add pane to title row.""" + lines = [line1, line2] + colors = colors if colors else self.title_colors.copy() + if not line2: + lines.pop() + colors.pop() tmux_args = { 'behind': True, 'lines': TMUX_TITLE_HEIGHT, 'target_id': None, - 'text': text, + 'text': ansi.color_string(lines, colors, sep='\n'), 'vertical': True, } if self.layout['Title']['Panes']: @@ -186,6 +191,7 @@ class TUI(): ), ) + # Functions def fix_layout(layout, forced=False): """Fix pane sizes based on layout.""" From 9678f143c70a78c9a3bf32a181914fa77f2605ec Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sun, 14 May 2023 21:59:34 -0700 Subject: [PATCH 12/16] Misc --- scripts/auto_repairs.py | 2 +- scripts/wk/exe.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/scripts/auto_repairs.py b/scripts/auto_repairs.py index b68f5e0b..aa0b4d29 100644 --- a/scripts/auto_repairs.py +++ b/scripts/auto_repairs.py @@ -5,7 +5,7 @@ import wk # Classes -REBOOT_STR = wk.ansi.color_string('Reboot', 'YELLOW') +REBOOT_STR = wk.ui.ansi.color_string('Reboot', 'YELLOW') class MenuEntry(): """Simple class to allow cleaner code below.""" def __init__(self, name, function=None, selected=True, **kwargs): diff --git a/scripts/wk/exe.py b/scripts/wk/exe.py index 5628754b..1a3ed084 100644 --- a/scripts/wk/exe.py +++ b/scripts/wk/exe.py @@ -1,5 +1,5 @@ """WizardKit: Execution functions""" -#vim: sts=2 sw=2 ts=2 +# vim: sts=2 sw=2 ts=2 import json import logging @@ -106,8 +106,8 @@ def build_cmd_kwargs(cmd, minimized=False, pipe=True, shell=False, **kwargs): # Start minimized if minimized: - startupinfo = subprocess.STARTUPINFO() - startupinfo.dwFlags = subprocess.STARTF_USESHOWWINDOW + startupinfo = subprocess.STARTUPINFO() # type: ignore + startupinfo.dwFlags = subprocess.STARTF_USESHOWWINDOW # type: ignore startupinfo.wShowWindow = 6 cmd_kwargs['startupinfo'] = startupinfo From ba69773fbad54a65be485721991c04af4a28c19e Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sat, 27 May 2023 19:08:48 -0700 Subject: [PATCH 13/16] Clear history when clearing a pane Helpful if the pane is resized to prevent cleared lines from returning. --- scripts/wk/ui/tmux.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/scripts/wk/ui/tmux.py b/scripts/wk/ui/tmux.py index d473f81f..4f376111 100644 --- a/scripts/wk/ui/tmux.py +++ b/scripts/wk/ui/tmux.py @@ -26,12 +26,16 @@ def capture_pane(pane_id=None): def clear_pane(pane_id=None): """Clear pane buffer for current or target pane.""" - cmd = ['tmux', 'send-keys', '-R'] + commands = [ + ['tmux', 'send-keys', '-R'], + ['tmux', 'clear-history'], + ] if pane_id: - cmd.extend(['-t', pane_id]) + commands = [[*cmd, '-t', pane_id] for cmd in commands] # Clear pane - run_program(cmd, check=False) + for cmd in commands: + run_program(cmd, check=False) def fix_layout(layout, forced=False): From 4c76e592385b7ddf9bd699d88d8ee83b2f37b2f6 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sat, 27 May 2023 19:22:24 -0700 Subject: [PATCH 14/16] Add get_window_size() --- scripts/wk/ui/tmux.py | 32 +++++++++++++++++++++++--------- 1 file changed, 23 insertions(+), 9 deletions(-) diff --git a/scripts/wk/ui/tmux.py b/scripts/wk/ui/tmux.py index 4f376111..780a6687 100644 --- a/scripts/wk/ui/tmux.py +++ b/scripts/wk/ui/tmux.py @@ -127,6 +127,20 @@ def get_pane_size(pane_id=None): return (width, height) +def get_window_size(): + """Get current window size, returns tuple.""" + cmd = ['tmux', 'display', '-p', '#{window_width} #{window_height}'] + + # Get resolution + proc = run_program(cmd, check=False) + width, height = proc.stdout.strip().split() + width = int(width) + height = int(height) + + # Done + return (width, height) + + def kill_all_panes(pane_id=None): """Kill all panes except for the current or target pane.""" cmd = ['tmux', 'kill-pane', '-a'] @@ -273,6 +287,15 @@ def resize_pane(pane_id=None, width=None, height=None, **kwargs): run_program(cmd, check=False) +def respawn_pane(pane_id, **action): + """Respawn pane with action.""" + cmd = ['tmux', 'respawn-pane', '-k', '-t', pane_id] + cmd.extend(prep_action(**action)) + + # Respawn + run_program(cmd, check=False) + + def split_window( lines=None, percent=None, behind=False, vertical=False, @@ -309,15 +332,6 @@ def split_window( return proc.stdout.strip() -def respawn_pane(pane_id, **action): - """Respawn pane with action.""" - cmd = ['tmux', 'respawn-pane', '-k', '-t', pane_id] - cmd.extend(prep_action(**action)) - - # Respawn - run_program(cmd, check=False) - - def zoom_pane(pane_id=None): """Toggle zoom status for current or target pane.""" cmd = ['tmux', 'resize-pane', '-Z'] From cb012423bb7d25ad9cfa9fdb0a1619ab52628a9d Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sat, 27 May 2023 19:47:26 -0700 Subject: [PATCH 15/16] Refactor hardware diagnostics to use new TUI --- scripts/wk/hw/cpu.py | 9 +-- scripts/wk/hw/diags.py | 165 ++++++++++++----------------------------- scripts/wk/ui/tmux.py | 50 ++++++------- scripts/wk/ui/tui.py | 117 ++++++++++++++++++++--------- 4 files changed, 158 insertions(+), 183 deletions(-) diff --git a/scripts/wk/hw/cpu.py b/scripts/wk/hw/cpu.py index 18a4bd83..597093ef 100644 --- a/scripts/wk/hw/cpu.py +++ b/scripts/wk/hw/cpu.py @@ -11,7 +11,7 @@ from wk import exe from wk.cfg.hw import CPU_FAILURE_TEMP from wk.os.mac import set_fans as macos_set_fans from wk.std import PLATFORM -from wk.ui import ansi, tmux +from wk.ui import ansi # STATIC VARIABLES @@ -111,7 +111,7 @@ def start_mprime(working_dir, log_path) -> subprocess.Popen: stdin=proc_mprime.stdout, stdout=subprocess.PIPE, ) - proc_mprime.stdout.close() + proc_mprime.stdout.close() # type: ignore[reportOptionalMemberAccess] save_nsbr = exe.NonBlockingStreamReader(proc_grep.stdout) exe.start_thread( save_nsbr.save_to_file, @@ -122,7 +122,7 @@ def start_mprime(working_dir, log_path) -> subprocess.Popen: return proc_mprime -def start_sysbench(sensors, sensors_out, log_path, pane) -> SysbenchType: +def start_sysbench(sensors, sensors_out, log_path) -> SysbenchType: """Start sysbench, returns tuple with Popen object and file handle.""" set_apple_fan_speed('max') sysbench_cmd = [ @@ -141,9 +141,6 @@ def start_sysbench(sensors, sensors_out, log_path, pane) -> SysbenchType: thermal_action=('killall', 'sysbench', '-INT'), ) - # Update bottom pane - tmux.respawn_pane(pane, watch_file=log_path, watch_cmd='tail') - # Start sysbench filehandle_sysbench = open( log_path, 'a', encoding='utf-8', diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py index 60e42abe..8b680a2d 100644 --- a/scripts/wk/hw/diags.py +++ b/scripts/wk/hw/diags.py @@ -25,8 +25,7 @@ from wk.hw.network import network_test from wk.hw.screensavers import screensaver from wk.hw.test import Test, TestGroup -from wk.ui import tui as ui -from wk.ui import ansi, cli, tmux +from wk.ui import ansi, cli, tui # STATIC VARIABLES @@ -85,33 +84,24 @@ class State(): def __init__(self, test_mode=False): self.disks = [] self.log_dir = None - self.panes = {} self.progress_file = None self.system = None self.test_groups = [] self.title_text = ansi.color_string('Hardware Diagnostics', 'GREEN') if test_mode: self.title_text += ansi.color_string(' (Test Mode)', 'YELLOW') - self.ui = ui.TUI(self.title_text) + self.ui = tui.TUI(f'{self.title_text}\nMain Menu') def abort_testing(self) -> None: - """Set unfinished tests as aborted and cleanup tmux panes.""" + """Set unfinished tests as aborted and cleanup panes.""" for group in self.test_groups: for test in group.test_objects: if test.status in ('Pending', 'Working'): test.set_status('Aborted') - # Cleanup tmux - self.panes.pop('Current', None) - for key, pane_ids in self.panes.copy().items(): - if key in ('Top', 'Started', 'Progress'): - continue - if isinstance(pane_ids, str): - tmux.kill_pane(self.panes.pop(key)) - else: - for _id in pane_ids: - tmux.kill_pane(_id) - self.panes.pop(key) + # Cleanup panes + self.ui.remove_all_info_panes() + self.ui.remove_all_worker_panes() def disk_safety_checks(self) -> None: """Check for mid-run SMART failures and failed test(s).""" @@ -127,8 +117,6 @@ class State(): # Reset objects self.disks.clear() - self.layout.clear() - self.layout.update(cfg.hw.TMUX_LAYOUT) self.test_groups.clear() # Set log @@ -191,35 +179,6 @@ class State(): test_group.test_objects.append(test_obj) self.test_groups.append(test_group) - def init_tmux(self) -> None: - """Initialize tmux layout.""" - tmux.kill_all_panes() - - # Top - self.panes['Top'] = tmux.split_window( - behind=True, - lines=2, - vertical=True, - text=f'{self.title_text}\nMain Menu', - ) - - # Started - self.panes['Started'] = tmux.split_window( - lines=cfg.hw.TMUX_SIDE_WIDTH, - target_id=self.panes['Top'], - text=ansi.color_string( - ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], - ['BLUE', None], - sep='\n', - ), - ) - - # Progress - self.panes['Progress'] = tmux.split_window( - lines=cfg.hw.TMUX_SIDE_WIDTH, - text=' ', - ) - def save_debug_reports(self) -> None: """Save debug reports to disk.""" LOG.info('Saving debug reports') @@ -266,17 +225,6 @@ class State(): _f.write(f'\n{test.name}:\n') _f.write('\n'.join(debug.generate_object_report(test, indent=1))) - def update_clock(self) -> None: - """Update 'Started' pane following clock sync.""" - tmux.respawn_pane( - pane_id=self.panes['Started'], - text=ansi.color_string( - ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], - ['BLUE', None], - sep='\n', - ), - ) - def update_progress_file(self) -> None: """Update progress file.""" report = [] @@ -296,9 +244,9 @@ class State(): # Write to progress file self.progress_file.write_text('\n'.join(report), encoding='utf-8') - def update_top_pane(self, text) -> None: + def update_title_text(self, text) -> None: """Update top pane with text.""" - tmux.respawn_pane(self.panes['Top'], text=f'{self.title_text}\n{text}') + self.ui.set_title(self.title_text, text) # Functions @@ -370,7 +318,7 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: return # Prep - state.update_top_pane(test_mprime_obj.dev.cpu_description) + state.update_title_text(test_mprime_obj.dev.cpu_description) test_cooling_obj.set_status('Working') test_mprime_obj.set_status('Working') @@ -383,17 +331,16 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: # Create monitor and worker panes state.update_progress_file() - state.panes['Prime95'] = tmux.split_window( - lines=10, vertical=True, watch_file=prime_log, watch_cmd='tail') + state.ui.add_worker_pane(lines=10, watch_cmd='tail', watch_file=prime_log) if PLATFORM == 'Darwin': - state.panes['Temps'] = tmux.split_window( - behind=True, percent=80, vertical=True, cmd='./hw-sensors') + state.ui.add_info_pane( + percent=80, cmd='./hw-sensors', update_layout=False, + ) elif PLATFORM == 'Linux': - state.panes['Temps'] = tmux.split_window( - behind=True, percent=80, vertical=True, watch_file=sensors_out) - tmux.resize_pane(height=3) - state.panes['Current'] = '' - state.layout['Current'] = {'height': 3, 'Check': True} + state.ui.add_info_pane( + percent=80, watch_file=sensors_out, update_layout=False, + ) + state.ui.set_current_pane_height(3) # Get idle temps cli.print_standard('Saving idle temps...') @@ -421,7 +368,7 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: state.update_progress_file() # Get cooldown temp - cli.clear_screen() + state.ui.clear_current_pane() cli.print_standard('Letting CPU cooldown...') std.sleep(5) cli.print_standard('Saving cooldown temps...') @@ -440,15 +387,18 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: if run_sysbench: LOG.info('CPU Test (Sysbench)') cli.print_standard('Letting CPU cooldown more...') - std.sleep(30) - cli.clear_screen() + std.sleep(10) + state.ui.clear_current_pane() cli.print_info('Running alternate stress test') print('') + sysbench_log = prime_log.with_name('sysbench.log') + sysbench_log.touch() + state.ui.remove_all_worker_panes() + state.ui.add_worker_pane(lines=10, watch_cmd='tail', watch_file=sysbench_log) proc_sysbench, filehandle_sysbench = hw_cpu.start_sysbench( sensors, sensors_out, - log_path=prime_log.with_name('sysbench.log'), - pane=state.panes['Prime95'], + log_path=sysbench_log, ) try: print_countdown(proc=proc_sysbench, seconds=test_minutes*60) @@ -474,9 +424,9 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: # Cleanup state.update_progress_file() sensors.stop_background_monitor() - state.panes.pop('Current', None) - tmux.kill_pane(state.panes.pop('Prime95', None)) - tmux.kill_pane(state.panes.pop('Temps', None)) + state.ui.clear_current_pane_height() + state.ui.remove_all_info_panes() + state.ui.remove_all_worker_panes() # Done if aborted: @@ -504,14 +454,10 @@ def disk_io_benchmark( aborted = False # Run benchmarks - state.update_top_pane( + state.update_title_text( f'Disk I/O Benchmark{"s" if len(test_objects) > 1 else ""}', ) - state.panes['I/O Benchmark'] = tmux.split_window( - percent=50, - vertical=True, - text=' ', - ) + state.ui.set_current_pane_height(10) for test in test_objects: if test.disabled: # Skip @@ -523,12 +469,14 @@ def disk_io_benchmark( continue # Start benchmark - cli.clear_screen() + state.ui.clear_current_pane() cli.print_report(test.dev.generate_report()) test.set_status('Working') test_log = f'{state.log_dir}/{test.dev.path.name}_benchmark.out' - tmux.respawn_pane( - state.panes['I/O Benchmark'], + state.ui.remove_all_worker_panes() + state.ui.add_worker_pane( + percent=50, + update_layout=False, watch_cmd='tail', watch_file=test_log, ) @@ -554,7 +502,8 @@ def disk_io_benchmark( # Cleanup state.update_progress_file() - tmux.kill_pane(state.panes.pop('I/O Benchmark', None)) + state.ui.clear_current_pane_height() + state.ui.remove_all_worker_panes() # Done if aborted: @@ -566,10 +515,9 @@ def disk_self_test(state, test_objects, test_mode=False) -> None: LOG.info('Disk Self-Test(s)') aborted = False threads = [] - state.panes['SMART'] = [] # Run self-tests - state.update_top_pane( + state.update_title_text( f'Disk self-test{"s" if len(test_objects) > 1 else ""}', ) cli.print_info(f'Starting self-test{"s" if len(test_objects) > 1 else ""}') @@ -586,9 +534,7 @@ def disk_self_test(state, test_objects, test_mode=False) -> None: # Show progress if threads[-1].is_alive(): - state.panes['SMART'].append( - tmux.split_window(lines=4, vertical=True, watch_file=test_log), - ) + state.ui.add_worker_pane(lines=4, watch_cmd='tail', watch_file=test_log) # Wait for all tests to complete state.update_progress_file() @@ -607,9 +553,7 @@ def disk_self_test(state, test_objects, test_mode=False) -> None: # Cleanup state.update_progress_file() - for pane in state.panes['SMART']: - tmux.kill_pane(pane) - state.panes.pop('SMART', None) + state.ui.remove_all_worker_panes() # Done if aborted: @@ -663,10 +607,9 @@ def disk_surface_scan(state, test_objects, test_mode=False) -> None: LOG.info('Disk Surface Scan (badblocks)') aborted = False threads = [] - state.panes['badblocks'] = [] # Update panes - state.update_top_pane( + state.update_title_text( f'Disk Surface Scan{"s" if len(test_objects) > 1 else ""}', ) cli.print_info( @@ -685,14 +628,7 @@ def disk_surface_scan(state, test_objects, test_mode=False) -> None: # Show progress if threads[-1].is_alive(): - state.panes['badblocks'].append( - tmux.split_window( - lines=5, - vertical=True, - watch_cmd='tail', - watch_file=test_log, - ), - ) + state.ui.add_worker_pane(lines=5, watch_cmd='tail', watch_file=test_log) # Wait for all tests to complete try: @@ -713,9 +649,7 @@ def disk_surface_scan(state, test_objects, test_mode=False) -> None: # Cleanup state.update_progress_file() - for pane in state.panes['badblocks']: - tmux.kill_pane(pane) - state.panes.pop('badblocks', None) + state.ui.remove_all_worker_panes() # Done if aborted: @@ -733,7 +667,6 @@ def main() -> None: raise RuntimeError('tmux session not found') # Init - atexit.register(tmux.kill_all_panes) menu = build_menu(cli_mode=args['--cli'], quick_mode=args['--quick']) state = State(test_mode=args['--test-mode']) @@ -759,7 +692,7 @@ def main() -> None: # Run simple test if action: - state.update_top_pane(selection[0]) + state.update_title_text(selection[0]) try: action() except KeyboardInterrupt: @@ -767,7 +700,7 @@ def main() -> None: cli.print_standard('') cli.pause('Press Enter to return to main menu...') if 'Clock Sync' in selection: - state.update_clock() + state.ui.update_clock() # Secrets if 'Matrix' in selection: @@ -791,7 +724,7 @@ def main() -> None: run_diags(state, menu, quick_mode=False, test_mode=args['--test-mode']) # Reset top pane - state.update_top_pane('Main Menu') + state.update_title_text('Main Menu') def print_countdown(proc, seconds) -> None: @@ -842,7 +775,7 @@ def run_diags(state, menu, quick_mode=False, test_mode=False) -> None: args = [group.test_objects] if group.name == 'Disk I/O Benchmark': args.append(menu.toggles['Skip USB Benchmarks']['Selected']) - cli.clear_screen() + state.ui.clear_current_pane() try: function(state, *args, test_mode=test_mode) except (KeyboardInterrupt, std.GenericAbort): @@ -887,8 +820,8 @@ def show_failed_attributes(state) -> None: def show_results(state) -> None: """Show test results by device.""" std.sleep(0.5) - cli.clear_screen() - state.update_top_pane('Results') + state.ui.clear_current_pane() + state.update_title_text('Results') # CPU Tests cpu_tests_enabled = [ diff --git a/scripts/wk/ui/tmux.py b/scripts/wk/ui/tmux.py index 780a6687..e653720d 100644 --- a/scripts/wk/ui/tmux.py +++ b/scripts/wk/ui/tmux.py @@ -51,13 +51,27 @@ def fix_layout(layout, forced=False): for data in layout.values(): data['Panes'] = [pane for pane in data['Panes'] if poll_pane(pane)] + # Calc height for "floating" row + # NOTE: We start with height +1 to account for the splits (i.e. splits = num rows - 1) + floating_height = 1 + get_window_size()[1] + for group in ('Title', 'Info', 'Current', 'Workers'): + if layout[group]['Panes']: + group_height = 1 + layout[group].get('height', 0) + if group == 'Workers': + group_height *= len(layout[group]['Panes']) + floating_height -= group_height + # Update main panes - resize_pane(height=999) # Set active pane too large and then adjust down for section, data in layout.items(): + # "Floating" pane(s) + if 'height' not in data and section in ('Info', 'Current', 'Workers'): + for pane_id in data['Panes']: + resize_kwargs.append({'pane_id': pane_id, 'height': floating_height}) + + # Rest of the panes if section == 'Workers': # Skip for now continue - if 'height' in data: for pane_id in data['Panes']: resize_kwargs.append({'pane_id': pane_id, 'height': data['height']}) @@ -81,33 +95,19 @@ def fix_layout(layout, forced=False): resize_pane(pane_id, width=width) if group == 'Title': # (re)fix Started pane - #TODO: REstore: resize_pane(layout['Started']['Panes'][0], width=TMUX_SIDE_WIDTH) - resize_pane(layout['Started']['Panes'][0], width=21) + resize_pane(layout['Started']['Panes'][0], width=layout['Started']['width']) # Bail early - if not (layout['Workers']['Panes'] and layout['Workers']['height']): + if not ( + layout['Workers']['Panes'] + and 'height' in layout['Workers'] + and floating_height > 0 + ): return # Update worker heights - worker_height = layout['Workers']['height'] - workers = layout['Workers']['Panes'].copy() - num_workers = len(workers) - avail_height = sum(get_pane_size(pane)[1] for pane in workers) - avail_height += get_pane_size()[1] # Current pane - # Check if window is too small - if avail_height < (worker_height*num_workers) + 3: - # Just leave things as-is - return - # Resize current pane - resize_pane(height=avail_height-(worker_height*num_workers)) - # Resize bottom pane - resize_pane(workers.pop(0), height=worker_height) - # Resize the rest of the panes by adjusting the ones above them - while len(workers) > 1: - next_height = sum(get_pane_size(pane)[1] for pane in workers[:2]) - next_height -= worker_height - resize_pane(workers[1], height=next_height) - workers.pop(0) + for worker in reversed(layout['Workers']['Panes']): + resize_pane(worker, height=layout['Workers']['height']) def get_pane_size(pane_id=None): @@ -262,7 +262,7 @@ def prep_file(path): pass -def resize_pane(pane_id=None, width=None, height=None, **kwargs): +def resize_pane(pane_id=None, width=None, height=None): """Resize current or target pane. NOTE: kwargs is only here to make calling this function easier diff --git a/scripts/wk/ui/tui.py b/scripts/wk/ui/tui.py index f7834701..95d86ae7 100644 --- a/scripts/wk/ui/tui.py +++ b/scripts/wk/ui/tui.py @@ -6,6 +6,7 @@ import logging import time from copy import deepcopy +from os import environ from wk.exe import start_thread from wk.std import sleep @@ -15,13 +16,13 @@ from wk.ui import ansi, tmux LOG = logging.getLogger(__name__) TMUX_SIDE_WIDTH = 21 TMUX_TITLE_HEIGHT = 2 -TMUX_LAYOUT = { - 'Current': {'Panes': [None]}, +TMUX_LAYOUT = { # NOTE: This needs to be in order from top to bottom 'Title': {'Panes': [], 'height': TMUX_TITLE_HEIGHT}, + 'Info': {'Panes': []}, + 'Current': {'Panes': [environ.get('TMUX_PANE', None)]}, + 'Workers': {'Panes': []}, 'Started': {'Panes': [], 'width': TMUX_SIDE_WIDTH}, 'Progress': {'Panes': [], 'width': TMUX_SIDE_WIDTH}, - 'Info': {'Panes': []}, - 'Workers': {'Panes': []}, } @@ -42,11 +43,22 @@ class TUI(): # Close all panes at exit atexit.register(tmux.kill_all_panes) - def add_info_pane(self, height, **tmux_args) -> None: + def add_info_pane( + self, lines=None, percent=None, update_layout=True, **tmux_args, + ) -> None: """Add info pane.""" + if not (lines or percent): + # Bail early + raise RuntimeError('Neither lines nor percent specified.') + + # Calculate lines if needed + if not lines: + lines = int(tmux.get_pane_size()[1] * (percent/100)) + + # Set tmux split args tmux_args.update({ 'behind': True, - 'lines': height, + 'lines': lines, 'target_id': None, 'vertical': True, }) @@ -60,7 +72,8 @@ class TUI(): tmux_args.pop('lines') # Update layout - self.layout['Info']['height'] = height + if update_layout: + self.layout['Info']['height'] = lines # Add pane self.layout['Info']['Panes'].append(tmux.split_window(**tmux_args)) @@ -91,19 +104,48 @@ class TUI(): # Add pane self.layout['Title']['Panes'].append(tmux.split_window(**tmux_args)) - def add_worker_pane(self, height, **tmux_split_args) -> None: + def add_worker_pane( + self, lines=None, percent=None, update_layout=True, **tmux_args, + ) -> None: """Add worker pane.""" - self.layout['Workers']['height'] = height - self.layout['Workers']['Panes'].append(tmux.split_window( - vertical=True, - lines=height, - **tmux_split_args, - )) + height = lines + + # Bail early + if not (lines or percent): + raise RuntimeError('Neither lines nor percent specified.') + + # Calculate height if needed + if not height: + height = int(tmux.get_pane_size()[1] * (percent/100)) + + # Set tmux split args + tmux_args.update({ + 'behind': False, + 'lines': lines, + 'percent': percent, + 'target_id': None, + 'vertical': True, + }) + + # Update layout + if update_layout: + self.layout['Workers']['height'] = height + + # Add pane + self.layout['Workers']['Panes'].append(tmux.split_window(**tmux_args)) + + def clear_current_pane(self) -> None: + """Clear screen and history for current pane.""" + tmux.clear_pane() + + def clear_current_pane_height(self) -> None: + """Clear current pane height and update layout.""" + self.layout['Current'].pop('height', None) def fix_layout(self, forced=True) -> None: """Fix tmux layout based on self.layout.""" try: - fix_layout(self.layout, forced=forced) + tmux.fix_layout(self.layout, forced=forced) except RuntimeError: # Assuming self.panes changed while running pass @@ -166,6 +208,11 @@ class TUI(): self.layout['Workers']['Panes'].clear() tmux.kill_pane(*panes) + def set_current_pane_height(self, height) -> None: + """Set current pane height and update layout.""" + self.layout['Current']['height'] = height + tmux.resize_pane(height=height) + def set_progress_file(self, progress_file) -> None: """Set the file to use for the progresse pane.""" tmux.respawn_pane( @@ -191,6 +238,17 @@ class TUI(): ), ) + def update_clock(self) -> None: + """Update 'Started' pane following clock sync.""" + tmux.respawn_pane( + pane_id=self.layout['Started']['Panes'][0], + text=ansi.color_string( + ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], + ['BLUE', None], + sep='\n', + ), + ) + # Functions def fix_layout(layout, forced=False): @@ -207,7 +265,6 @@ def fix_layout(layout, forced=False): data['Panes'] = [pane for pane in data['Panes'] if tmux.poll_pane(pane)] # Update main panes - tmux.resize_pane(height=999) # Set active pane too large and then adjust down for section, data in layout.items(): if section == 'Workers': # Skip for now @@ -278,31 +335,19 @@ def layout_needs_fixed(layout): tmux.get_pane_size(pane)[0] != data['width'] for pane in data['Panes'] ) - # TODO: Re-enable? - ## Group panes - #for group in ('Title', 'Info'): - # num_panes = len(layout[group]['Panes']) - # if num_panes <= 1: - # continue - # width = int( (tmux.get_pane_size()[0] - (1 - num_panes)) / num_panes ) - # for pane in layout[group]['Panes']: - # needs_fixed = needs_fixed or abs(tmux.get_pane_size(pane)[0] - width) > 2 - # Done return needs_fixed -def hmm(): - """Hmm?""" - - def test(): + """TODO: Deleteme""" ui = TUI() - ui.add_info_pane(height=10, text='Info One') - ui.add_info_pane(height=10, text='Info Two') - ui.add_info_pane(height=10, text='Info Three') - ui.add_worker_pane(height=3, text='Work One') - ui.add_worker_pane(height=3, text='Work Two') - ui.add_worker_pane(height=3, text='Work Three') + ui.add_info_pane(lines=10, text='Info One') + ui.add_info_pane(lines=10, text='Info Two') + ui.add_info_pane(lines=10, text='Info Three') + ui.add_worker_pane(lines=3, text='Work One') + ui.add_worker_pane(lines=3, text='Work Two') + ui.add_worker_pane(lines=3, text='Work Three') + ui.fix_layout() return ui From 3334638a2c87fb81f4d2e3956f1a9f5b5b08ae5d Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sat, 27 May 2023 19:48:18 -0700 Subject: [PATCH 16/16] Suppress warnings when using tail in tmux --- scripts/wk/ui/tmux.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/wk/ui/tmux.py b/scripts/wk/ui/tmux.py index e653720d..22438c3e 100644 --- a/scripts/wk/ui/tmux.py +++ b/scripts/wk/ui/tmux.py @@ -242,7 +242,7 @@ def prep_action( 'cat', ]) elif watch_cmd == 'tail': - action_cmd.extend(['tail', '-f']) + action_cmd.extend(['tail', '-q', '-f']) action_cmd.append(watch_file) else: LOG.error('No action specified')