diff --git a/scripts/auto_repairs.py b/scripts/auto_repairs.py index c01f1735..aa0b4d29 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.ansi.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..973b76e2 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.ansi.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..81188020 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.ansi.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..afb3e8d5 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 @@ -17,21 +17,22 @@ from . import os 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/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/clone/ddrescue.py b/scripts/wk/clone/ddrescue.py index 29667bd9..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 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, @@ -33,6 +33,8 @@ from wk.hw.smart import ( smart_status_ok, update_smart_details, ) +from wk.ui import cli as ui +from wk.ui import ansi, tmux # STATIC VARIABLES @@ -89,8 +91,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 {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, @@ -296,7 +298,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 +422,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=ansi.color_string( ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], ['BLUE', None], sep='\n', @@ -442,7 +444,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 +456,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 +490,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,9 +524,9 @@ class State(): # New run, use new settings file settings['Needs Format'] = True offset = 0 - user_choice = std.choice( - ['G', 'M', 'S'], + user_choice = ui.choice( 'Format clone using GPT, MBR, or match Source type?', + ['G', 'M', 'S'], ) if user_choice == 'G': settings['Table Type'] = 'GPT' @@ -533,7 +535,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 @@ -561,19 +563,19 @@ 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 = [] # Source - report.append(std.color_string('Source', 'GREEN')) + report.append(ansi.color_string('Source', 'GREEN')) report.extend(build_object_report(self.source)) report.append(' ') # Destination - report.append(std.color_string('Destination', 'GREEN')) + report.append(ansi.color_string('Destination', 'GREEN')) if self.mode == 'Clone': - report[-1] += std.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(' ') @@ -581,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(std.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( - std.color_string( + ansi.color_string( ['This is irreversible and will lead to', 'DATA LOSS.'], ['YELLOW', 'RED'], ), @@ -605,18 +607,18 @@ class State(): # Map dir if self.working_dir: - report.append(std.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( - std.color_string( + ansi.color_string( 'Map file(s) are being saved to a non-recommended filesystem.', 'YELLOW', ), ) report.append( - std.color_string( + ansi.color_string( ['This is strongly discouraged and may lead to', 'DATA LOSS'], [None, 'RED'], ), @@ -625,11 +627,11 @@ class State(): # Source part(s) selected if source_parts: - report.append(std.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(std.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} ' @@ -638,9 +640,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_msg): raise std.GenericAbort() def generate_report(self): @@ -661,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 = std.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}') @@ -673,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 = std.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) @@ -704,7 +706,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 @@ -744,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, ) @@ -795,8 +797,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 +930,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 +971,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 +1028,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 +1039,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 +1066,10 @@ class State(): width = cfg.ddrescue.TMUX_SIDE_WIDTH # Status - report.append(std.color_string(f'{"Status":^{width}}', 'BLUE')) + report.append(ansi.color_string(f'{"Status":^{width}}', 'BLUE')) if 'NEEDS ATTENTION' in overall_status: report.append( - std.color_string(f'{overall_status:^{width}}', 'YELLOW_BLINK'), + ansi.color_string(f'{overall_status:^{width}}', 'YELLOW_BLINK'), ) else: report.append(f'{overall_status:^{width}}') @@ -1077,12 +1079,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(ansi.color_string('Overall Progress', 'BLUE')) report.append( f'Rescued: {format_status_string(percent, width=width-9)}', ) report.append( - std.color_string( + ansi.color_string( [f'{std.bytes_to_string(total_rescued, decimals=2):>{width}}'], [get_percent_color(percent)], ), @@ -1091,7 +1093,7 @@ class State(): # Block pair progress for pair in self.block_pairs: - report.append(std.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( @@ -1103,9 +1105,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(ansi.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(ansi.color_string('N/A', 'YELLOW')) else: report.append(etoc) @@ -1166,7 +1168,7 @@ class State(): source_str = _format_string(self.source, width) tmux.respawn_pane( self.panes['Source'], - text=std.color_string( + text=ansi.color_string( ['Source', '' if source_exists else ' (Missing)', '\n', source_str], ['BLUE', 'RED', None, None], sep='', @@ -1181,7 +1183,7 @@ class State(): percent=50, vertical=False, target_id=self.panes['Source'], - text=std.color_string( + text=ansi.color_string( ['Destination', '' if dest_exists else ' (Missing)', '\n', dest_str], ['BLUE', 'RED', None, None], sep='', @@ -1195,7 +1197,7 @@ def build_block_pair_report(block_pairs, settings): report = [] notes = [] if block_pairs: - report.append(std.color_string('Block Pairs', 'GREEN')) + report.append(ansi.color_string('Block Pairs', 'GREEN')) else: # Bail early return report @@ -1214,7 +1216,7 @@ def build_block_pair_report(block_pairs, settings): if settings: if not settings['First Run']: notes.append( - std.color_string( + ansi.color_string( ['NOTE:', 'Clone settings loaded from previous run.'], ['BLUE', None], ), @@ -1222,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( - std.color_string( + ansi.color_string( ['NOTE:', msg], ['BLUE', None], ), ) if any(pair.get_rescued_size() > 0 for pair in block_pairs): notes.append( - std.color_string( + ansi.color_string( ['NOTE:', 'Resume data loaded from map file(s).'], ['BLUE', None], ), @@ -1311,12 +1313,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 = ansi.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(ansi.color_string('PATH', 'BLUE')) report.append(str(path)) # Done @@ -1352,7 +1354,7 @@ def build_disk_report(dev): # Partition details report.append( - std.color_string( + ansi.color_string( ( f'{"NAME":<{widths["name"]}}' f'{" " if dev.children else ""}' @@ -1397,8 +1399,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=ansi.color_string('ddrescue TUI: Main Menu', 'GREEN')) menu.separator = ' ' # Add actions, options, etc @@ -1429,23 +1431,25 @@ 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'), + ansi.color_string('ddrescue TUI: Expert Settings', 'GREEN'), ' ', - std.color_string( + ansi.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:') + ui.print_standard( + f'Available ddrescue presets: {" / ".join(SETTING_PRESETS)}' + ) + preset = ui.choice('Please select a preset:', SETTING_PRESETS) # Fix selection for _p in SETTING_PRESETS: @@ -1494,7 +1498,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 +1580,7 @@ def format_status_string(status, width): # Add color if necessary if color: - status_str = std.color_string(status_str, color) + status_str = ansi.color_string(status_str, color) # Done return status_str @@ -1698,8 +1702,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 +1714,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 +1779,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 @@ -1784,30 +1788,20 @@ 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 = std.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: 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 +1845,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 +1881,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 +1904,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 +1922,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 +1943,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', ansi.strip_colors(line)) def mount_raw_image(path): @@ -1970,7 +1964,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 +2031,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 +2050,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 +2061,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 +2074,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( + ansi.color_string( ['SMART Attributes', f'Updated: {now}\n'], ['BLUE', 'YELLOW'], sep='\t\t', @@ -2113,7 +2107,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: @@ -2159,19 +2153,19 @@ 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(' ') - std.print_error('DDRESCUE PROCESS HALTED') - print(' ') - std.print_warning(warning_message) + ui.print_standard(' ') + ui.print_standard(' ') + ui.print_error('DDRESCUE PROCESS HALTED') + ui.print_standard(' ') + 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 +2181,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 +2242,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() @@ -2258,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().""" - 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=ansi.color_string(f'ddrescue TUI: {prompt_msg} Selection', 'GREEN'), ) menu.disabled_str = 'Already selected' menu.separator = ' ' @@ -2303,11 +2297,11 @@ 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 = std.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 = std.Menu(title) + menu = ui.Menu(title) menu.separator = ' ' menu.add_action('All') menu.add_action('None') @@ -2319,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(): @@ -2356,7 +2350,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 += ansi.color_string(' No partitions detected.', 'YELLOW') # Get selection _select_parts(menu) @@ -2369,25 +2363,25 @@ 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()}?' - if std.ask(msg): + 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] # Convert object_list to hw_disk.Disk() objects - print(' ') - std.print_info('Getting disk/partition details...') + ui.print_standard(' ') + ui.print_info('Getting disk/partition details...') object_list = [hw_disk.Disk(path) for path in object_list] # Done return object_list -def select_path(prompt): +def select_path(prompt_msg): """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=ansi.color_string(f'ddrescue TUI: {prompt_msg} Path Selection', 'GREEN'), ) menu.separator = ' ' menu.add_action('Quit') @@ -2400,7 +2394,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 +2404,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 +2423,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('Are we cloning or imaging?', ['C', 'I']) 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/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 diff --git a/scripts/wk/graph.py b/scripts/wk/graph.py index 885cfe8e..95324956 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 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 = color_string((GRAPH_HORIZONTAL[-1],), (rate_color,)) + full_block = ansi.color_string((GRAPH_HORIZONTAL[-1],), (rate_color,)) if step >= 24: - graph[0] += 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] += 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] += 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] += 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 = 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 8aebf51d..71a0a465 100644 --- a/scripts/wk/hw/benchmark.py +++ b/scripts/wk/hw/benchmark.py @@ -21,11 +21,8 @@ from wk.cfg.hw import ( THRESH_SSD_MIN, ) from wk.exe import run_program -from wk.std import ( - PLATFORM, - strip_colors, - color_string, - ) +from wk.std import PLATFORM +from wk.ui import ansi # 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 ansi.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(ansi.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'), + 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 8458b918..597093ef 100644 --- a/scripts/wk/hw/cpu.py +++ b/scripts/wk/hw/cpu.py @@ -10,13 +10,8 @@ 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, - color_string, - print_error, - print_warning, - ) -from wk.tmux import respawn_pane as tmux_respawn_pane +from wk.std import PLATFORM +from wk.ui import ansi # STATIC VARIABLES @@ -70,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( @@ -97,9 +92,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(ansi.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(ansi.color_string(' Unknown result', 'YELLOW')) def start_mprime(working_dir, log_path) -> subprocess.Popen: @@ -116,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, @@ -127,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 = [ @@ -146,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', @@ -174,9 +166,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}') - for line in str(err).splitlines(): - 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/diags.py b/scripts/wk/hw/diags.py index 7dc44870..999451d1 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 @@ -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 ansi, cli, tui + # STATIC VARIABLES DOCSTRING = f'''{cfg.main.KIT_NAME_FULL}: Hardware Diagnostics @@ -81,37 +83,25 @@ 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 = std.color_string('Hardware Diagnostics', 'GREEN') + self.title_text = ansi.color_string('Hardware Diagnostics', 'GREEN') if test_mode: - self.top_text += std.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 = 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).""" @@ -128,30 +118,11 @@ 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.""" # Reset objects self.disks.clear() - self.layout.clear() - self.layout.update(cfg.hw.TMUX_LAYOUT) self.test_groups.clear() # Set log @@ -166,15 +137,13 @@ class State(): keep_history=False, timestamp=False, ) - std.clear_screen() - std.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() @@ -216,35 +185,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.top_text}\nMain Menu', - ) - - # Started - self.panes['Started'] = tmux.split_window( - lines=cfg.hw.TMUX_SIDE_WIDTH, - target_id=self.panes['Top'], - text=std.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') @@ -253,7 +193,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))) @@ -291,27 +231,15 @@ 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=std.color_string( - ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], - ['BLUE', None], - sep='\n', - ), - ) - - 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(std.color_string(group.name, 'BLUE')) + report.append(ansi.color_string(group.name, 'BLUE')) for test in group.test_objects: - report.append(std.color_string( - [test.label, f'{test.status:>{width-len(test.label)}}'], + report.append(ansi.color_string( + [test.label, f'{test.status:>{self.ui.side_width-len(test.label)}}'], [None, STATUS_COLORS.get(test.status, None)], sep='', )) @@ -320,19 +248,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: + def update_title_text(self, text) -> None: """Update top pane with text.""" - tmux.respawn_pane(self.panes['Top'], text=f'{self.top_text}\n{text}') + self.ui.set_title(self.title_text, text) # 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) -> cli.Menu: + """Build main menu, returns wk.ui.cli.Menu.""" + menu = cli.Menu(title=None) # Add actions, options, etc for action in MENU_ACTIONS: @@ -398,7 +324,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') @@ -410,25 +336,24 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: ) # Create monitor and worker panes - state.update_progress_pane() - state.panes['Prime95'] = tmux.split_window( - lines=10, vertical=True, watch_file=prime_log, watch_cmd='tail') + state.update_progress_file() + 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 - std.print_standard('Saving idle temps...') + cli.print_standard('Saving idle temps...') sensors.save_average_temps(temp_label='Idle', seconds=5) # Stress CPU - std.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) @@ -446,17 +371,17 @@ 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 - std.clear_screen() - std.print_standard('Letting CPU cooldown...') + state.ui.clear_current_pane() + cli.print_standard('Letting CPU cooldown...') std.sleep(5) - std.print_standard('Saving cooldown temps...') + cli.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(ansi.color_string('Prime95', 'BLUE')) hw_cpu.check_mprime_results( test_obj=test_mprime_obj, working_dir=state.log_dir, ) @@ -467,16 +392,19 @@ 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...') - std.sleep(30) - std.clear_screen() - std.print_info('Running alternate stress test') + cli.print_standard('Letting CPU cooldown more...') + 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) @@ -493,18 +421,18 @@ 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(std.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 - 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)) - 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: @@ -522,7 +450,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( @@ -532,14 +460,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 @@ -551,16 +475,18 @@ def disk_io_benchmark( continue # Start benchmark - std.clear_screen() - std.print_report(test.dev.generate_report()) + 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, ) - state.update_progress_pane() + state.update_progress_file() try: hw_benchmark.run_io_test(test, test_log, test_mode=test_mode) except KeyboardInterrupt: @@ -569,20 +495,21 @@ 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(ansi.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(ansi.color_string(' Aborted', 'YELLOW')) break # Update progress after each test - state.update_progress_pane() + state.update_progress_file() # Cleanup - state.update_progress_pane() - tmux.kill_pane(state.panes.pop('I/O Benchmark', None)) + state.update_progress_file() + state.ui.clear_current_pane_height() + state.ui.remove_all_worker_panes() # Done if aborted: @@ -594,13 +521,12 @@ 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 ""}', ) - std.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: @@ -614,12 +540,10 @@ 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_pane() + state.update_progress_file() try: while True: if any(t.is_alive() for t in threads): @@ -634,10 +558,8 @@ 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() - for pane in state.panes['SMART']: - tmux.kill_pane(pane) - state.panes.pop('SMART', None) + state.update_progress_file() + state.ui.remove_all_worker_panes() # Done if aborted: @@ -691,13 +613,12 @@ 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 ""}', ) - std.print_info( + cli.print_info( f'Starting disk surface scan{"s" if len(test_objects) > 1 else ""}', ) show_failed_attributes(state) @@ -713,20 +634,13 @@ 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: 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,13 +651,11 @@ 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(ansi.color_string(' Aborted', 'YELLOW')) # Cleanup - state.update_progress_pane() - for pane in state.panes['badblocks']: - tmux.kill_pane(pane) - state.panes.pop('badblocks', None) + state.update_progress_file() + state.ui.remove_all_worker_panes() # Done if aborted: @@ -761,7 +673,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']) @@ -787,15 +698,15 @@ def main() -> None: # Run simple test if action: - state.update_top_pane(selection[0]) + state.update_title_text(selection[0]) try: action() except KeyboardInterrupt: - std.print_warning('Aborted.') - std.print_standard('') - std.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() + state.ui.update_clock() # Secrets if 'Matrix' in selection: @@ -819,7 +730,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: @@ -858,8 +769,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() + cli.print_warning('No tests selected?') + cli.pause() return # Run tests @@ -870,13 +781,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']) - std.clear_screen() + state.ui.clear_current_pane() 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 @@ -897,48 +808,48 @@ 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...') + cli.pause('Press Enter to exit...') else: - std.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: - std.print_colored([dev.name, dev.description], ['CYAN', None]) - std.print_report( + cli.print_colored([dev.name, dev.description], ['CYAN', None]) + cli.print_report( hw_smart.generate_attribute_report(dev, only_failed=True), ) - std.print_standard('') + cli.print_standard('') def show_results(state) -> None: """Show test results by device.""" std.sleep(0.5) - std.clear_screen() - state.update_top_pane('Results') + state.ui.clear_current_pane() + state.update_title_text('Results') # CPU Tests cpu_tests_enabled = [ 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(' ') + 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: - std.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: - std.print_report(disk.generate_report()) - std.print_standard(' ') + cli.print_report(disk.generate_report()) + cli.print_standard(' ') if not state.disks: - std.print_warning('No devices') - std.print_standard(' ') + cli.print_warning('No devices') + cli.print_standard(' ') def sync_clock() -> None: diff --git a/scripts/wk/hw/disk.py b/scripts/wk/hw/disk.py index 3c7c727c..31704c8f 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 import ansi # STATIC VARIABLES @@ -73,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 = ansi.color_string(note, color) if note not in self.notes: self.notes.append(note) self.notes.sort() @@ -82,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 == ansi.strip_colors(note): present = True return present @@ -98,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(ansi.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(ansi.color_string('Attributes', 'BLUE')) report.extend(generate_attribute_report(self)) # Notes if self.notes: - report.append(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/keyboard.py b/scripts/wk/hw/keyboard.py index 68e2d0a6..07d3a22d 100644 --- a/scripts/wk/hw/keyboard.py +++ b/scripts/wk/hw/keyboard.py @@ -4,7 +4,7 @@ import logging from wk.exe import run_program -from wk.std import PLATFORM, print_warning +from wk.std import PLATFORM # STATIC VARIABLES @@ -17,7 +17,8 @@ def keyboard_test() -> None: if PLATFORM == 'Linux': run_xev() else: - 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: diff --git a/scripts/wk/hw/network.py b/scripts/wk/hw/network.py index 700ebfea..55d02eb2 100644 --- a/scripts/wk/hw/network.py +++ b/scripts/wk/hw/network.py @@ -9,11 +9,7 @@ from wk.net import ( show_valid_addresses, speedtest, ) -from wk.std import ( - TryAndPrint, - pause, - print_warning, - ) +from wk.ui import cli as ui # STATIC VARIABLES @@ -24,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, @@ -34,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) @@ -51,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/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/hw/sensors.py b/scripts/wk/hw/sensors.py index 49d693f0..815a6de5 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 import ansi # STATIC VARIABLES @@ -109,7 +110,7 @@ class Sensors(): # Handle empty reports if not report: report = [ - color_string('WARNING: No sensors found', 'YELLOW'), + ansi.color_string('WARNING: No sensors found', 'YELLOW'), '', 'Please monitor temps manually', ] @@ -425,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 ansi.color_string(temp, 'PURPLE') # Determine color if colored: @@ -435,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 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 9bce00f9..11163f9b 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 import ansi # STATIC VARIABLES @@ -40,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 = [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(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(color_string(f' {test_obj.status}', 'YELLOW')) + report.append(ansi.color_string(f' {test_obj.status}', 'YELLOW')) elif test_obj.status == 'TestInProgress': - report.append(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(color_string(' Aborted', 'YELLOW')) + report.append(ansi.color_string(' Aborted', 'YELLOW')) test_obj.set_status('Aborted') elif test_obj.status == 'TimedOut': - report.append(color_string(' TimedOut', 'YELLOW')) + report.append(ansi.color_string(' TimedOut', 'YELLOW')) # Done test_obj.report.extend(report) @@ -136,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 = ansi.color_string( [label, get_attribute_value_string(dev, attr), note], [None, value_color, 'YELLOW'], ) @@ -324,7 +325,7 @@ def run_smart_self_test(test_obj, log_path) -> None: 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 = 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 49afd25c..135c9951 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 import ansi # STATIC VARIABLES @@ -31,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 = ansi.strip_colors(line.strip()) if not line or BADBLOCKS_SKIP_REGEX.match(line): # Skip continue @@ -48,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' {ansi.color_string(line, "YELLOW")}') test_obj.set_status('Failed') else: - test_obj.report.append(f' {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') @@ -65,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(ansi.color_string('badblocks', 'BLUE')) test_obj.set_status('Working') # Increase block size if necessary @@ -84,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( + 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 8da2ed15..4309c3df 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 import ansi # STATIC VARIABLES @@ -41,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(ansi.color_string('Device', 'BLUE')) report.append(f' {self.cpu_description}') # Include RAM details - report.append(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/kit/build_win.py b/scripts/wk/kit/build_win.py index ff377790..e87d5101 100644 --- a/scripts/wk/kit/build_win.py +++ b/scripts/wk/kit/build_win.py @@ -21,16 +21,8 @@ from wk.kit.tools import ( get_tool_path, ) from wk.log import update_log_path -from wk.std import ( - GenericError, - TryAndPrint, - clear_screen, - pause, - print_info, - print_success, - set_title, - sleep, - ) +from wk.std import GenericError +from wk.ui import cli as ui # STATIC VARIABLES @@ -245,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(): @@ -315,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) @@ -459,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'): @@ -496,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, @@ -514,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 cc5e9969..64697a41 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 + # 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..5e478e54 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 import cli as ui # REGEX @@ -199,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 481e687c..ee17792a 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 import ansi # STATIC VARIABLES @@ -82,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 = ansi.color_string( [f'{vol["name"]:<20}', f'{vol["size"]:>9}'], [None, 'CYAN'], ) # Mountpoint and type - line = 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 = 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 a5043d3d..4d6d7248 100644 --- a/scripts/wk/os/win.py +++ b/scripts/wk/os/win.py @@ -30,10 +30,10 @@ from wk.std import ( GenericError, GenericWarning, bytes_to_string, - color_string, - input_text, sleep, ) +from wk.ui import cli as ui +from wk.ui import ansi # STATIC VARIABLES @@ -183,7 +183,7 @@ def check_4k_alignment(show_alert=False): continue if int(match.group('offset')) % 4096 != 0: report.append( - color_string( + ansi.color_string( f'{match.group("description")}' f' ({bytes_to_string(match.group("size"), decimals=1)})' , @@ -199,7 +199,7 @@ def check_4k_alignment(show_alert=False): if report: report.insert( 0, - color_string('One or more partitions not 4K aligned', 'YELLOW'), + ansi.color_string('One or more partitions not 4K aligned', 'YELLOW'), ) return report @@ -212,7 +212,7 @@ def export_bitlocker_info(): ] # Get filename - file_name = 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 @@ -251,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(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(color_string('No products detected', 'RED')) + report.append(ansi.color_string('No products detected', 'RED')) # Done return report @@ -364,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 = 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 938fb3b3..51c418d4 100644 --- a/scripts/wk/repairs/win.py +++ b/scripts/wk/repairs/win.py @@ -58,21 +58,10 @@ from wk.os.win import ( from wk.std import ( GenericError, GenericWarning, - Menu, - TryAndPrint, - abort, - ask, - clear_screen, - color_string, - pause, - print_info, - print_standard, - print_warning, - set_title, - show_data, sleep, - strip_colors, ) +from wk.ui import cli as ui +from wk.ui import ansi # STATIC VARIABLES @@ -87,7 +76,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( @@ -105,7 +94,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'): @@ -116,7 +105,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{ansi.color_string("Main Menu", "GREEN")}') # Main Menu for entry in base_menus['Actions']: @@ -125,7 +114,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{ansi.color_string("Options", "GREEN")}') for entry in base_menus['Options']: menus['Options'].add_option(entry.name, entry.details) menus['Options'].add_action('All') @@ -135,7 +124,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{ansi.color_string(group, "GREEN")}') menus[group].disabled_str = 'Locked' for entry in entries: menus[group].add_option(entry.name, entry.details) @@ -167,7 +156,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{ansi.color_string("Load Preset", "GREEN")}' MENU_PRESETS.add_option('Default') for name in presets: MENU_PRESETS.add_option(name) @@ -271,7 +260,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) @@ -316,7 +305,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 @@ -392,20 +381,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, ansi.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 @@ -423,21 +412,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(): @@ -446,19 +435,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 = ansi.strip_colors(name) skipped = details.get('Skipped', False) done = details.get('Done', False) disabled = details.get('Disabled', False) @@ -472,7 +461,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 @@ -482,7 +471,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, ) @@ -490,7 +479,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 @@ -513,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}\{strip_colors(name)}' + key_path = fr'{AUTO_REPAIR_KEY}\{group}\{ansi.strip_colors(name)}' # Get values from TryAndPrint result if result: @@ -527,7 +516,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 = ansi.strip_colors(value_name) if isinstance(data, bool): data = 1 if data else 0 if isinstance(data, int): @@ -891,7 +880,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) @@ -905,7 +894,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, ) @@ -933,7 +922,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, ) @@ -1319,7 +1308,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 2e5cfbf8..97f93b5b 100644 --- a/scripts/wk/setup/win.py +++ b/scripts/wk/setup/win.py @@ -60,22 +60,10 @@ from wk.repairs.win import ( from wk.std import ( GenericError, GenericWarning, - Menu, - TryAndPrint, - abort, - ask, - clear_screen, - color_string, - pause, - print_error, - print_info, - print_standard, - print_warning, - set_title, - show_data, sleep, - strip_colors, ) +from wk.ui import cli as ui +from wk.ui import ansi # STATIC VARIABLES @@ -91,7 +79,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)', @@ -103,7 +91,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'): @@ -114,7 +102,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{ansi.color_string("Main Menu", "GREEN")}') # Main Menu for entry in base_menus['Actions']: @@ -124,7 +112,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{ansi.color_string(group, "GREEN")}') for entry in entries: menus[group].add_option(entry.name, entry.details) menus[group].add_action('All') @@ -153,7 +141,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{ansi.color_string("Load Preset", "GREEN")}' MENU_PRESETS.add_option('Default') for name in presets: MENU_PRESETS.add_option(name) @@ -170,26 +158,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} ({ansi.color_string(os_name, color)})' def load_preset(menus, presets, title, enable_menu_exit=True): @@ -215,10 +203,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 @@ -235,11 +223,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) @@ -254,10 +242,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(): @@ -266,22 +254,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 = ansi.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 @@ -418,7 +406,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) @@ -804,8 +792,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() @@ -906,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(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/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..75a335fc --- /dev/null +++ b/scripts/wk/ui/__init__.py @@ -0,0 +1,6 @@ +"""WizardKit: ui module init""" + +from . import ansi +from . import cli +from . import tmux +from . import tui diff --git a/scripts/wk/ui/ansi.py b/scripts/wk/ui/ansi.py new file mode 100644 index 00000000..2c21c0ed --- /dev/null +++ b/scripts/wk/ui/ansi.py @@ -0,0 +1,72 @@ +"""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 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'] + 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/ui/cli.py b/scripts/wk/ui/cli.py new file mode 100644 index 00000000..3fb7c173 --- /dev/null +++ b/scripts/wk/ui/cli.py @@ -0,0 +1,880 @@ +"""WizardKit: CLI functions""" +# vim: sts=2 sw=2 ts=2 + +import logging +import os +import platform +import re +import subprocess +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 +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) +from wk.ui.ansi import clear_screen, color_string, strip_colors + +# STATIC VARIABLES +LOG = logging.getLogger(__name__) +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. + + 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_msg): + """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_msg).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_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_msg) + 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_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'], + } + + while True: + self._update(single_selection=True, settings_mode=True) + 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': + # 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_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_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. + + 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_msg='Aborted.', show_prompt_msg=True, return_code=1): + """Abort script.""" + print_warning(prompt_msg) + if show_prompt_msg: + sleep(0.5) + pause(prompt_msg='Press Enter to exit... ') + sys.exit(return_code) + + +def ask(prompt_msg): + """Prompt the user with a Y/N question, returns bool.""" + validator = InputYesNoValidator() + + # 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_msg, '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(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('prompt_msg: %s, choices: %s', prompt_msg, choices) + choices = [str(c).upper()[:1] for c in choices] + prompt_msg = f'{prompt_msg} [{"/".join(choices)}]' + + # Show prompt + response = input_text(prompt_msg, validator=InputChoiceValidator(choices)) + + # Done + 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 +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_ticket_id(): + """Get ticket ID, returns str.""" + prompt_msg = 'Please enter ticket ID:' + validator = InputTicketIDValidator() + + # 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: + result = prompt(prompt_msg, validator=validator) + except KeyboardInterrupt: + # Ignore CTRL+c + pass + + # Done + return result + + +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_msg='Press Enter to continue... '): + """Simple pause implementation.""" + input_text(prompt_msg, allow_empty=True) + + +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='', + ) + + +if __name__ == '__main__': + print("This file is not meant to be called directly.") diff --git a/scripts/wk/tmux.py b/scripts/wk/ui/tmux.py similarity index 60% rename from scripts/wk/tmux.py rename to scripts/wk/ui/tmux.py index 49fb3241..22438c3e 100644 --- a/scripts/wk/tmux.py +++ b/scripts/wk/ui/tmux.py @@ -26,38 +26,88 @@ 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(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: - continue + # Remove closed panes + for data in layout.values(): + data['Panes'] = [pane for pane in data['Panes'] if poll_pane(pane)] - # 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 + # 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 + 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']}) + 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 + resize_pane(layout['Started']['Panes'][0], width=layout['Started']['width']) + + # Bail early + if not ( + layout['Workers']['Panes'] + and 'height' in layout['Workers'] + and floating_height > 0 + ): + return + + # Update worker heights + for worker in reversed(layout['Workers']['Panes']): + resize_pane(worker, height=layout['Workers']['height']) def get_pane_size(pane_id=None): @@ -77,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'] @@ -96,34 +160,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 @@ -182,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') @@ -202,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 @@ -227,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, @@ -263,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'] diff --git a/scripts/wk/ui/tui.py b/scripts/wk/ui/tui.py new file mode 100644 index 00000000..95d86ae7 --- /dev/null +++ b/scripts/wk/ui/tui.py @@ -0,0 +1,355 @@ +"""WizardKit: TUI functions""" +# vim: sts=2 sw=2 ts=2 + +import atexit +import logging +import time + +from copy import deepcopy +from os import environ + +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 = { # 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}, + } + + +# 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) + + # Close all panes at exit + atexit.register(tmux.kill_all_panes) + + 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': lines, + '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 + if update_layout: + self.layout['Info']['height'] = lines + + # Add pane + self.layout['Info']['Panes'].append(tmux.split_window(**tmux_args)) + + 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': ansi.color_string(lines, colors, sep='\n'), + '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, lines=None, percent=None, update_layout=True, **tmux_args, + ) -> None: + """Add worker pane.""" + 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: + tmux.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_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( + 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', + ), + ) + + 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): + """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 + 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'] + ) + + # Done + return needs_fixed + +def test(): + """TODO: Deleteme""" + ui = TUI() + 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 + + +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() 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