From e13b057eac46130878277849c10ecc7dbda39b0f Mon Sep 17 00:00:00 2001 From: Alan Mason <1923621+2Shirt@users.noreply.github.com> Date: Fri, 1 Dec 2017 19:24:09 -0800 Subject: [PATCH] Sync functions\* with WinPE versions --- .bin/Scripts/functions/common.py | 37 +++++++++++++--- .bin/Scripts/functions/data.py | 76 ++++++++++++++++---------------- 2 files changed, 70 insertions(+), 43 deletions(-) diff --git a/.bin/Scripts/functions/common.py b/.bin/Scripts/functions/common.py index ad4a3c49..285129c7 100644 --- a/.bin/Scripts/functions/common.py +++ b/.bin/Scripts/functions/common.py @@ -37,6 +37,9 @@ class BIOSKeyNotFoundError(Exception): class BinNotFoundError(Exception): pass +class GenericAbort(Exception): + pass + class GenericError(Exception): pass @@ -221,6 +224,10 @@ def menu_select(title='~ Untitled Menu ~', if not main_entries and not action_entries: raise Exception("MenuError: No items given") + # Set title + if 'Title' in global_vars: + title = '{}\n\n{}'.format(global_vars['Title'], title) + # Build menu menu_splash = '{}\n\n'.format(title) width = len(str(len(main_entries))) @@ -263,7 +270,7 @@ def menu_select(title='~ Untitled Menu ~', answer = '' while (answer.upper() not in valid_answers): - clear_screen() + os.system('cls') print(menu_splash) answer = input('{}: '.format(prompt)) @@ -358,10 +365,29 @@ def run_program(cmd, args=[], check=True, pipe=True, shell=False): return process_return -def show_info(message='~Some message~', info='~Some info~', indent=8, width=32): +def set_title(title='~Some Title~'): + """Set title. + + Used for window title and menu titles.""" + global_vars['Title'] = title + os.system('title {}'.format(title)) + +def show_data(message='~Some message~', data='~Some data~', indent=8, width=32, + info=False, warning=False, error=False): """Display info with formatting.""" - print_standard('{indent}{message:<{width}}{info}'.format( - indent=' '*indent, width=width, message=message, info=info)) + message = '{indent}{message:<{width}}{data}'.format( + indent=' '*indent, width=width, message=message, data=data) + if error: + print_error(message) + elif warning: + print_warning(message) + elif info: + print_info(message) + else: + print_standard(message) + +def show_info(message='~Some message~', info='~Some info~', indent=8, width=32): + show_data(message=message, data=info, indent=indent, width=width) def sleep(seconds=2): """Wait for a while.""" @@ -400,6 +426,7 @@ def try_and_print(message='Trying...', and the result string will be printed in the correct color. catch_all=False will result in unspecified exceptions being re-raised.""" err = None + out = None w_exceptions = other_results.get('Warning', {}).keys() w_exceptions = tuple(get_exception(e) for e in w_exceptions) e_exceptions = other_results.get('Error', {}).keys() @@ -435,7 +462,7 @@ def try_and_print(message='Trying...', if err and not catch_all: raise else: - return {'CS': not bool(err), 'Error': err} + return {'CS': not bool(err), 'Error': err, 'Out': out} def upload_data(path, file): """Add CLIENT_INFO_SERVER to authorized connections and upload file.""" diff --git a/.bin/Scripts/functions/data.py b/.bin/Scripts/functions/data.py index 12d568b7..f61b51c0 100644 --- a/.bin/Scripts/functions/data.py +++ b/.bin/Scripts/functions/data.py @@ -92,7 +92,7 @@ FAST_COPY_EXCLUDES = [ r'Temporary?Items', r'Thumbs.db', ] -FAST_COPY_ARGS = [ +FAST_COPY_ARGS = [ '/cmd=noexist_only', '/utf8', '/skip_empty_dir', @@ -388,43 +388,6 @@ def select_destination(folder_path, prompt='Select destination'): return path -def select_volume(title='Select disk', auto_select=True): - """Select disk from attached disks. returns dict.""" - actions = [{'Name': 'Quit', 'Letter': 'Q'}] - disks = [] - - # Build list of disks - set_thread_error_mode(silent=True) # Prevents "No disk" popups - for d in psutil.disk_partitions(): - info = { - 'Disk': d, - 'Name': d.mountpoint} - try: - usage = psutil.disk_usage(d.device) - free = '{free} / {total} available'.format( - free = human_readable_size(usage.free, 2), - total = human_readable_size(usage.total, 2)) - except Exception: - # Meh, leaving unsupported destinations out - pass - # free = 'Unknown' - # info['Disabled'] = True - else: - info['Display Name'] = '{} ({})'.format(info['Name'], free) - disks.append(info) - set_thread_error_mode(silent=False) # Return to normal - - # Skip menu? - if len(disks) == 1 and auto_select: - return disks[0] - - # Show menu - selection = menu_select(title, main_entries=disks, action_entries=actions) - if selection == 'Q': - exit_script() - else: - return disks[int(selection)-1] - def select_source(ticket_number): """Select backup from those found on the BACKUP_SERVERS for the ticket.""" selected_source = None @@ -536,6 +499,43 @@ def select_source(ticket_number): # Done return selected_source +def select_volume(title='Select disk', auto_select=True): + """Select disk from attached disks. returns dict.""" + actions = [{'Name': 'Quit', 'Letter': 'Q'}] + disks = [] + + # Build list of disks + set_thread_error_mode(silent=True) # Prevents "No disk" popups + for d in psutil.disk_partitions(): + info = { + 'Disk': d, + 'Name': d.mountpoint} + try: + usage = psutil.disk_usage(d.device) + free = '{free} / {total} available'.format( + free = human_readable_size(usage.free, 2), + total = human_readable_size(usage.total, 2)) + except Exception: + # Meh, leaving unsupported destinations out + pass + # free = 'Unknown' + # info['Disabled'] = True + else: + info['Display Name'] = '{} ({})'.format(info['Name'], free) + disks.append(info) + set_thread_error_mode(silent=False) # Return to normal + + # Skip menu? + if len(disks) == 1 and auto_select: + return disks[0] + + # Show menu + selection = menu_select(title, main_entries=disks, action_entries=actions) + if selection == 'Q': + exit_script() + else: + return disks[int(selection)-1] + def set_thread_error_mode(silent=True): """Disable or Enable Windows error message dialogs.