From c009ab2d4181c1a44b21de7fee7bde17b6e2e97d Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Mon, 29 May 2023 16:04:58 -0700 Subject: [PATCH] Add even more type hints to function arguments --- scripts/wk/exe.py | 4 +- scripts/wk/ui/cli.py | 138 ++++++++++++++++++++++++++---------------- scripts/wk/ui/tmux.py | 47 ++++++++------ scripts/wk/ui/tui.py | 36 ++++++++--- 4 files changed, 145 insertions(+), 80 deletions(-) diff --git a/scripts/wk/exe.py b/scripts/wk/exe.py index 00ab62d6..2fb3ff98 100644 --- a/scripts/wk/exe.py +++ b/scripts/wk/exe.py @@ -212,7 +212,7 @@ def popen_program( minimized: bool = False, pipe: bool = False, shell: bool = False, - **kwargs: dict[Any, Any], + **kwargs, ) -> subprocess.Popen: """Run program and return a subprocess.Popen object.""" LOG.debug( @@ -242,7 +242,7 @@ def run_program( check: bool = True, pipe: bool = True, shell: bool = False, - **kwargs: dict[Any, Any], + **kwargs, ) -> subprocess.CompletedProcess: """Run program and return a subprocess.CompletedProcess object.""" LOG.debug( diff --git a/scripts/wk/ui/cli.py b/scripts/wk/ui/cli.py index 14051884..8a2eaf46 100644 --- a/scripts/wk/ui/cli.py +++ b/scripts/wk/ui/cli.py @@ -10,9 +10,10 @@ import sys import traceback from collections import OrderedDict -from typing import Any +from typing import Any, Callable, Iterable from prompt_toolkit import prompt +from prompt_toolkit.document import Document from prompt_toolkit.validation import Validator, ValidationError try: @@ -38,12 +39,12 @@ PLATFORM = platform.system() # Classes class InputChoiceValidator(Validator): """Validate that input is one of the provided choices.""" - def __init__(self, choices, allow_empty=False): + def __init__(self, choices: Iterable[str], allow_empty: bool = False): self.allow_empty = allow_empty self.choices = [str(c).upper() for c in choices] super().__init__() - def validate(self, document) -> None: + def validate(self, document: Document) -> None: text = document.text if not (text or self.allow_empty): raise ValidationError( @@ -58,7 +59,7 @@ class InputChoiceValidator(Validator): class InputNotEmptyValidator(Validator): """Validate that input is not empty.""" - def validate(self, document) -> None: + def validate(self, document: Document) -> None: text = document.text if not text: raise ValidationError( @@ -68,11 +69,11 @@ class InputNotEmptyValidator(Validator): class InputTicketIDValidator(Validator): """Validate that input resembles a ticket ID.""" - def __init__(self, allow_empty=False): + def __init__(self, allow_empty: bool = False): self.allow_empty = allow_empty super().__init__() - def validate(self, document) -> None: + def validate(self, document: Document) -> None: text = document.text if not (text or self.allow_empty): raise ValidationError( @@ -87,11 +88,11 @@ class InputTicketIDValidator(Validator): class InputYesNoValidator(Validator): """Validate that input is a yes or no.""" - def __init__(self, allow_empty=False): + def __init__(self, allow_empty: bool = False): self.allow_empty = allow_empty super().__init__() - def validate(self, document) -> None: + def validate(self, document: Document) -> None: text = document.text if not (text or self.allow_empty): raise ValidationError( @@ -113,7 +114,7 @@ class Menu(): 1. All entry names are unique. 2. All action entry names start with different letters. """ - def __init__(self, title='[Untitled Menu]'): + def __init__(self, title: str = '[Untitled Menu]'): self.actions = OrderedDict() self.options = OrderedDict() self.sets = OrderedDict() @@ -236,7 +237,7 @@ class Menu(): # Done return valid_answers - def _resolve_selection(self, selection) -> tuple[str, dict[Any, Any]]: + def _resolve_selection(self, selection: str) -> tuple[str, dict[Any, Any]]: """Get menu item based on user selection, returns tuple.""" offset = 1 resolved_selection = tuple() @@ -267,7 +268,7 @@ class Menu(): # Done return resolved_selection - def _update(self, single_selection=True, settings_mode=False) -> None: + def _update(self, single_selection: bool = True, settings_mode: bool = False) -> None: """Update menu items in preparation for printing to screen.""" index = 0 @@ -305,7 +306,8 @@ class Menu(): no_checkboxes=True, ) - def _update_entry_selection_status(self, entry, toggle=True, status=None) -> None: + def _update_entry_selection_status( + self, entry: str, toggle: bool = True, status: bool = False) -> None: """Update entry selection status either directly or by toggling.""" if entry in self.sets: # Update targets not the set itself @@ -319,14 +321,14 @@ class Menu(): else: section[entry]['Selected'] = status - def _update_set_selection_status(self, targets, status) -> None: + def _update_set_selection_status(self, targets: Iterable[str], status: bool) -> None: """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) -> str: + def _user_select(self, prompt_msg: str) -> str: """Show menu and select an entry, returns str.""" menu_text = self._generate_menu_text() valid_answers = self._get_valid_answers() @@ -343,19 +345,19 @@ class Menu(): # Done return answer - def add_action(self, name, details=None) -> None: + def add_action(self, name: str, details: dict[Any, Any] | None = None) -> 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) -> None: + def add_option(self, name: str, details: dict[Any, Any] | None = None) -> 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) -> None: + def add_set(self, name: str, details: dict[Any, Any] | None = None) -> None: """Add set to menu.""" details = details if details else {} details['Selected'] = details.get('Selected', False) @@ -367,14 +369,16 @@ class Menu(): # Add set self.sets[name] = details - def add_toggle(self, name, details=None) -> None: + def add_toggle(self, name: str, details: dict[Any, Any] | None = None) -> 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: ') -> tuple[str, dict[Any, Any]]: + self, + prompt_msg: str = 'Please make a selection: ', + ) -> tuple[str, dict[Any, Any]]: """Display menu and make multiple selections, returns tuple. NOTE: Menu is displayed until an action entry is selected. @@ -394,7 +398,9 @@ class Menu(): return selected_entry def settings_select( - self, prompt_msg='Please make a selection: ') -> tuple[str, dict[Any, Any]]: + self, + prompt_msg: str = 'Please make a selection: ', + ) -> tuple[str, dict[Any, Any]]: """Display menu and make multiple selections, returns tuple. NOTE: Menu is displayed until an action entry is selected. @@ -423,8 +429,10 @@ class Menu(): return selected_entry def simple_select( - self, prompt_msg='Please make a selection: ', - update=True) -> tuple[str, dict[Any, Any]]: + self, + prompt_msg: str = 'Please make a selection: ', + update: bool = True, + ) -> tuple[str, dict[Any, Any]]: """Display menu and make a single selection, returns tuple.""" if update: self._update() @@ -441,7 +449,7 @@ class TryAndPrint(): 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'): + def __init__(self, msg_bad: str = 'FAILED', msg_good: str = 'SUCCESS'): self.catch_all = True self.indent = INDENT self.list_errors = ['GenericError'] @@ -451,7 +459,7 @@ class TryAndPrint(): self.verbose = False self.width = WIDTH - def _format_exception_message(self, _exception) -> str: + def _format_exception_message(self, _exception: Exception) -> str: """Format using the exception's args or name, returns str.""" LOG.debug( 'Formatting exception: %s, %s', @@ -498,7 +506,11 @@ class TryAndPrint(): # Done return message - def _format_function_output(self, output, msg_good) -> str: + def _format_function_output( + self, + output: list | subprocess.CompletedProcess, + msg_good: str, + ) -> str: """Format function output for use in try_and_print(), returns str.""" LOG.debug('Formatting output: %s', output) @@ -536,27 +548,33 @@ class TryAndPrint(): # Done return result_msg - def _log_result(self, message, result_msg) -> None: + def _log_result(self, message: str, result_msg: str) -> None: """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) -> None: + def add_error(self, exception_name: str) -> None: """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) -> None: + def add_warning(self, exception_name: str) -> None: """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) -> dict[str, Any]: + self, + message: str, + function: Callable, + *args: Iterable[Any], + catch_all: bool | None = None, + msg_good: str | None = None, + verbose: bool | None = None, + **kwargs, + ) -> dict[str, Any]: """Run a function and print the results, returns results as dict. If catch_all is True then (nearly) all exceptions will be caught. @@ -594,8 +612,8 @@ class TryAndPrint(): 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) + e_exceptions: tuple = tuple(get_exception(e) for e in self.list_errors) + w_exceptions: tuple = 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) @@ -644,7 +662,10 @@ class TryAndPrint(): # Functions def abort( - prompt_msg='Aborted.', show_prompt_msg=True, return_code=1) -> None: + prompt_msg: str = 'Aborted.', + show_prompt_msg: bool = True, + return_code: int = 1, + ) -> None: """Abort script.""" print_warning(prompt_msg) if show_prompt_msg: @@ -653,7 +674,7 @@ def abort( sys.exit(return_code) -def ask(prompt_msg) -> bool: +def ask(prompt_msg: str) -> bool: """Prompt the user with a Y/N question, returns bool.""" validator = InputYesNoValidator() @@ -670,7 +691,7 @@ def ask(prompt_msg) -> bool: raise ValueError(f'Invalid answer given: {response}') -def beep(repeat=1) -> None: +def beep(repeat: int = 1) -> None: """Play system bell with optional repeat.""" while repeat >= 1: # Print bell char without a newline @@ -679,7 +700,7 @@ def beep(repeat=1) -> None: repeat -= 1 -def choice(prompt_msg, choices) -> str: +def choice(prompt_msg: str, choices: Iterable[str]) -> str: """Choose an option from a provided list, returns str. Choices provided will be converted to uppercase and returned as such. @@ -697,7 +718,7 @@ def choice(prompt_msg, choices) -> str: return response.upper() -def fix_prompt(message) -> str: +def fix_prompt(message: str) -> str: """Fix prompt, returns str.""" if not message: message = 'Input text: ' @@ -708,7 +729,7 @@ def fix_prompt(message) -> str: @cache -def get_exception(name) -> Exception: +def get_exception(name: str) -> Exception: """Get exception by name, returns exception object. [Doctest] @@ -757,7 +778,9 @@ def get_ticket_id() -> str: def input_text( - prompt_msg='Enter text: ', allow_empty=False, validator=None, + prompt_msg: str = 'Enter text: ', + allow_empty: bool = False, + validator: Validator | None = None, ) -> str: """Get input from user, returns str.""" prompt_msg = fix_prompt(prompt_msg) @@ -793,12 +816,18 @@ def major_exception() -> None: raise SystemExit(1) -def pause(prompt_msg='Press Enter to continue... ') -> None: +def pause(prompt_msg: str = 'Press Enter to continue... ') -> None: """Simple pause implementation.""" input_text(prompt_msg, allow_empty=True) -def print_colored(strings, colors, log=False, sep=' ', **kwargs) -> None: +def print_colored( + strings: Iterable[str] | str, + colors: Iterable[str | None] | str, + log: bool = False, + sep: str = ' ', + **kwargs, + ) -> None: """Prints strings in the colors specified.""" LOG.debug( 'strings: %s, colors: %s, sep: %s, kwargs: %s', @@ -816,7 +845,7 @@ def print_colored(strings, colors, log=False, sep=' ', **kwargs) -> None: LOG.info(strip_colors(msg)) -def print_error(msg, log=True, **kwargs) -> None: +def print_error(msg: str, log: bool = True, **kwargs) -> None: """Prints message in RED and log as ERROR.""" if 'file' not in kwargs: # Only set if not specified @@ -826,14 +855,14 @@ def print_error(msg, log=True, **kwargs) -> None: LOG.error(msg) -def print_info(msg, log=True, **kwargs) -> None: +def print_info(msg: str, log: bool = True, **kwargs) -> None: """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) -> None: +def print_report(report: list[str], indent=None, log: bool = True) -> None: """Print report to screen and optionally to log.""" for line in report: if indent: @@ -843,21 +872,21 @@ def print_report(report, indent=None, log=True) -> None: LOG.info(strip_colors(line)) -def print_standard(msg, log=True, **kwargs) -> None: +def print_standard(msg: str, log: bool = True, **kwargs) -> None: """Prints message and log as INFO.""" print(msg, **kwargs) if log: LOG.info(msg) -def print_success(msg, log=True, **kwargs) -> None: +def print_success(msg: str, log: bool = True, **kwargs) -> None: """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) -> None: +def print_warning(msg: str, log: bool = True, **kwargs) -> None: """Prints message in YELLOW and log as WARNING.""" if 'file' not in kwargs: # Only set if not specified @@ -867,7 +896,7 @@ def print_warning(msg, log=True, **kwargs) -> None: LOG.warning(msg) -def set_title(title) -> None: +def set_title(title: str) -> None: """Set window title.""" LOG.debug('title: %s', title) if os.name == 'nt': @@ -876,14 +905,19 @@ def set_title(title) -> None: print_error('Setting the title is only supported under Windows.') -def show_data(message, data, color=None, indent=None, width=None) -> None: +def show_data( + message: str, + data: Any, + color: str | None = None, + indent: int | None = None, + width: int | None = None, + ) -> 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, + (None, color if color else None), log=True, sep='', ) diff --git a/scripts/wk/ui/tmux.py b/scripts/wk/ui/tmux.py index bcd12759..ca040514 100644 --- a/scripts/wk/ui/tmux.py +++ b/scripts/wk/ui/tmux.py @@ -4,6 +4,8 @@ import logging import pathlib +from typing import Any + from wk.exe import run_program from wk.std import PLATFORM @@ -13,7 +15,7 @@ LOG = logging.getLogger(__name__) # Functions -def capture_pane(pane_id=None) -> str: +def capture_pane(pane_id: str | None = None) -> str: """Capture text from current or target pane, returns str.""" cmd = ['tmux', 'capture-pane', '-p'] if pane_id: @@ -24,7 +26,7 @@ def capture_pane(pane_id=None) -> str: return proc.stdout.strip() -def clear_pane(pane_id=None) -> None: +def clear_pane(pane_id: str | None = None) -> None: """Clear pane buffer for current or target pane.""" commands = [ ['tmux', 'send-keys', '-R'], @@ -38,7 +40,7 @@ def clear_pane(pane_id=None) -> None: run_program(cmd, check=False) -def fix_layout(layout, forced=False) -> None: +def fix_layout(layout: dict[str, dict[str, Any]], forced: bool = False) -> None: """Fix pane sizes based on layout.""" resize_kwargs = [] @@ -110,7 +112,7 @@ def fix_layout(layout, forced=False) -> None: resize_pane(worker, height=layout['Workers']['height']) -def get_pane_size(pane_id=None) -> tuple[int, int]: +def get_pane_size(pane_id: str | None = None) -> tuple[int, int]: """Get current or target pane size, returns tuple.""" cmd = ['tmux', 'display', '-p'] if pane_id: @@ -141,7 +143,7 @@ def get_window_size() -> tuple[int, int]: return (width, height) -def kill_all_panes(pane_id=None) -> None: +def kill_all_panes(pane_id: str | None = None) -> None: """Kill all panes except for the current or target pane.""" cmd = ['tmux', 'kill-pane', '-a'] if pane_id: @@ -151,7 +153,7 @@ def kill_all_panes(pane_id=None) -> None: run_program(cmd, check=False) -def kill_pane(*pane_ids) -> None: +def kill_pane(*pane_ids: str) -> None: """Kill pane(s) by id.""" cmd = ['tmux', 'kill-pane', '-t'] @@ -160,7 +162,7 @@ def kill_pane(*pane_ids) -> None: run_program(cmd+[pane_id], check=False) -def layout_needs_fixed(layout) -> bool: +def layout_needs_fixed(layout: dict[str, dict[str, Any]]) -> bool: """Check if layout needs fixed, returns bool.""" needs_fixed = False @@ -189,7 +191,7 @@ def layout_needs_fixed(layout) -> bool: return needs_fixed -def poll_pane(pane_id) -> bool: +def poll_pane(pane_id: str) -> bool: """Check if pane exists, returns bool.""" cmd = ['tmux', 'list-panes', '-F', '#D'] @@ -202,8 +204,12 @@ def poll_pane(pane_id) -> bool: def prep_action( - cmd=None, working_dir=None, text=None, - watch_file=None, watch_cmd='cat') -> list[str]: + cmd: str | None = None, + working_dir: pathlib.Path | str | None = None, + text: str | None = None, + watch_file: pathlib.Path | str | None = None, + watch_cmd: str = 'cat', + ) -> list[str]: """Prep action to perform during a tmux call, returns list. This will prep for running a basic command, displaying text on screen, @@ -253,7 +259,7 @@ def prep_action( return action_cmd -def prep_file(path) -> None: +def prep_file(path: pathlib.Path | str) -> None: """Check if file exists and create empty file if not.""" path = pathlib.Path(path).resolve() try: @@ -263,7 +269,11 @@ def prep_file(path) -> None: pass -def resize_pane(pane_id=None, width=None, height=None) -> None: +def resize_pane( + pane_id: str | None = None, + width: int | None = None, + height: int | None = None, + ) -> None: """Resize current or target pane. NOTE: kwargs is only here to make calling this function easier @@ -288,7 +298,7 @@ def resize_pane(pane_id=None, width=None, height=None) -> None: run_program(cmd, check=False) -def respawn_pane(pane_id, **action) -> None: +def respawn_pane(pane_id: str, **action) -> None: """Respawn pane with action.""" cmd = ['tmux', 'respawn-pane', '-k', '-t', pane_id] cmd.extend(prep_action(**action)) @@ -298,9 +308,12 @@ def respawn_pane(pane_id, **action) -> None: def split_window( - lines=None, percent=None, - behind=False, vertical=False, - target_id=None, **action) -> str: + lines: int | None = None, + percent: int | None = None, + behind: bool = False, + vertical: bool = False, + target_id: str | None = None, + **action) -> str: """Split tmux window, run action, and return pane_id as str.""" cmd = ['tmux', 'split-window', '-d', '-PF', '#D'] @@ -333,7 +346,7 @@ def split_window( return proc.stdout.strip() -def zoom_pane(pane_id=None) -> None: +def zoom_pane(pane_id: str | None = None) -> None: """Toggle zoom status for current or target pane.""" cmd = ['tmux', 'resize-pane', '-Z'] if pane_id: diff --git a/scripts/wk/ui/tui.py b/scripts/wk/ui/tui.py index ac6329b4..c5a68209 100644 --- a/scripts/wk/ui/tui.py +++ b/scripts/wk/ui/tui.py @@ -29,7 +29,7 @@ TMUX_LAYOUT = { # NOTE: This needs to be in order from top to bottom # Classes class TUI(): """Object for tracking TUI elements.""" - def __init__(self, title_text=None): + def __init__(self, title_text: str | None = None): self.layout = deepcopy(TMUX_LAYOUT) self.side_width = TMUX_SIDE_WIDTH self.title_text = title_text if title_text else 'Title Text' @@ -44,7 +44,11 @@ class TUI(): atexit.register(tmux.kill_all_panes) def add_info_pane( - self, lines=None, percent=None, update_layout=True, **tmux_args, + self, + lines: int | None = None, + percent: int | None = None, + update_layout: bool = True, + **tmux_args, ) -> None: """Add info pane.""" if not (lines or percent): @@ -78,7 +82,12 @@ class TUI(): # Add pane self.layout['Info']['Panes'].append(tmux.split_window(**tmux_args)) - def add_title_pane(self, line1, line2=None, colors=None) -> None: + def add_title_pane( + self, + line1: str, + line2: str | None = None, + colors: list[str] | None = None, + ) -> None: """Add pane to title row.""" lines = [line1, line2] colors = colors if colors else self.title_colors.copy() @@ -105,7 +114,11 @@ class TUI(): self.layout['Title']['Panes'].append(tmux.split_window(**tmux_args)) def add_worker_pane( - self, lines=None, percent=None, update_layout=True, **tmux_args, + self, + lines: int | None = None, + percent: int | None = None, + update_layout: bool = True, + **tmux_args, ) -> None: """Add worker pane.""" height = lines @@ -142,7 +155,7 @@ class TUI(): """Clear current pane height and update layout.""" self.layout['Current'].pop('height', None) - def fix_layout(self, forced=True) -> None: + def fix_layout(self, forced: bool = True) -> None: """Fix tmux layout based on self.layout.""" try: tmux.fix_layout(self.layout, forced=forced) @@ -208,19 +221,24 @@ class TUI(): self.layout['Workers']['Panes'].clear() tmux.kill_pane(*panes) - def set_current_pane_height(self, height) -> None: + def set_current_pane_height(self, height: int) -> 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: + def set_progress_file(self, progress_file: str) -> 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: + def set_title( + self, + line1: str, + line2: str | None = None, + colors: list[str] | None = None, + ) -> None: """Set title text.""" self.title_text = line1 self.title_text_line2 = line2 if line2 else '' @@ -251,7 +269,7 @@ class TUI(): # Functions -def fix_layout(layout, forced=False) -> None: +def fix_layout(layout, forced: bool = False) -> None: """Fix pane sizes based on layout.""" resize_kwargs = []