"""WizardKit: TUI functions""" # vim: sts=2 sw=2 ts=2 import logging import time from copy import deepcopy from wk.exe import start_thread from wk.std import sleep from wk.ui import ansi, tmux # STATIC VARIABLES LOG = logging.getLogger(__name__) TMUX_SIDE_WIDTH = 21 TMUX_TITLE_HEIGHT = 2 TMUX_LAYOUT = { 'Current': {'Panes': [None]}, 'Title': {'Panes': [], 'height': TMUX_TITLE_HEIGHT}, 'Started': {'Panes': [], 'width': TMUX_SIDE_WIDTH}, 'Progress': {'Panes': [], 'width': TMUX_SIDE_WIDTH}, 'Info': {'Panes': []}, 'Workers': {'Panes': []}, } # Classes class TUI(): """Object for tracking TUI elements.""" def __init__(self, title_text=None) -> None: self.layout = deepcopy(TMUX_LAYOUT) self.side_width = TMUX_SIDE_WIDTH self.title_text = title_text if title_text else 'Title Text' self.title_text_line2 = '' self.title_colors = ['BLUE', None] # Init tmux and start a background process to maintain layout self.init_tmux() start_thread(self.fix_layout_loop) def add_info_pane(self, height, **tmux_args) -> None: """Add info pane.""" tmux_args.update({ 'behind': True, 'lines': height, 'target_id': None, 'vertical': True, }) if self.layout['Info']['Panes']: tmux_args.update({ 'behind': False, 'percent': 50, 'target_id': self.layout['Info']['Panes'][-1], 'vertical': False, }) tmux_args.pop('lines') # Update layout self.layout['Info']['height'] = height # Add pane self.layout['Info']['Panes'].append(tmux.split_window(**tmux_args)) def add_title_pane(self, text) -> None: """Add additional pane to title row.""" tmux_args = { 'behind': True, 'lines': TMUX_TITLE_HEIGHT, 'target_id': None, 'text': text, 'vertical': True, } if self.layout['Title']['Panes']: tmux_args.update({ 'behind': False, 'percent': 50, 'target_id': self.layout['Title']['Panes'][-1], 'vertical': False, }) tmux_args.pop('lines') # Add pane self.layout['Title']['Panes'].append(tmux.split_window(**tmux_args)) def add_worker_pane(self, height, **tmux_split_args) -> None: """Add worker pane.""" self.layout['Workers']['height'] = height self.layout['Workers']['Panes'].append(tmux.split_window( vertical=True, lines=height, **tmux_split_args, )) def fix_layout(self, forced=True) -> None: """Fix tmux layout based on self.layout.""" try: fix_layout(self.layout, forced=forced) except RuntimeError: # Assuming self.panes changed while running pass def fix_layout_loop(self) -> None: """Fix layout on a loop. NOTE: This should be called as a thread. """ while True: self.fix_layout(forced=False) sleep(1) def init_tmux(self) -> None: """Initialize tmux layout.""" tmux.kill_all_panes() self.layout.clear() self.layout.update(deepcopy(TMUX_LAYOUT)) # Title self.layout['Title']['Panes'].append(tmux.split_window( behind=True, lines=2, vertical=True, text=ansi.color_string( [self.title_text, self.title_text_line2], self.title_colors, sep = '\n', ), )) # Started self.layout['Started']['Panes'].append(tmux.split_window( lines=TMUX_SIDE_WIDTH, target_id=self.layout['Title']['Panes'][0], text=ansi.color_string( ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], ['BLUE', None], sep='\n', ), )) # Progress self.layout['Progress']['Panes'].append(tmux.split_window( lines=TMUX_SIDE_WIDTH, text=' ', )) def remove_all_info_panes(self) -> None: """Remove all info panes and update layout.""" self.layout['Info'].pop('height', None) panes = self.layout['Info']['Panes'].copy() self.layout['Info']['Panes'].clear() tmux.kill_pane(*panes) def remove_all_worker_panes(self) -> None: """Remove all worker panes and update layout.""" self.layout['Workers'].pop('height', None) panes = self.layout['Workers']['Panes'].copy() self.layout['Workers']['Panes'].clear() tmux.kill_pane(*panes) def set_progress_file(self, progress_file) -> None: """Set the file to use for the progresse pane.""" tmux.respawn_pane( pane_id=self.layout['Progress']['Panes'][0], watch_file=progress_file, ) def set_title(self, line1, line2=None, colors=None) -> None: """Set title text.""" self.title_text = line1 self.title_text_line2 = line2 if line2 else '' if colors: self.title_colors = colors # Update pane (if present) if self.layout['Title']['Panes']: tmux.respawn_pane( pane_id=self.layout['Title']['Panes'][0], text=ansi.color_string( [self.title_text, self.title_text_line2], self.title_colors, sep = '\n', ), ) # Functions def fix_layout(layout, forced=False): """Fix pane sizes based on layout.""" resize_kwargs = [] # Bail early if not (forced or layout_needs_fixed(layout)): # Layout should be fine return # Remove closed panes for data in layout.values(): data['Panes'] = [pane for pane in data['Panes'] if tmux.poll_pane(pane)] # Update main panes tmux.resize_pane(height=999) # Set active pane too large and then adjust down for section, data in layout.items(): if section == 'Workers': # Skip for now continue if 'height' in data: for pane_id in data['Panes']: resize_kwargs.append({'pane_id': pane_id, 'height': data['height']}) if 'width' in data: for pane_id in data['Panes']: resize_kwargs.append({'pane_id': pane_id, 'width': data['width']}) for kwargs in resize_kwargs: try: tmux.resize_pane(**kwargs) except RuntimeError: # Assuming pane was closed just before resizing pass # Update "group" panes widths for group in ('Title', 'Info'): num_panes = len(layout[group]['Panes']) if num_panes <= 1: continue width = int( (tmux.get_pane_size()[0] - (1 - num_panes)) / num_panes ) for pane_id in layout[group]['Panes']: tmux.resize_pane(pane_id, width=width) if group == 'Title': # (re)fix Started pane tmux.resize_pane(layout['Started']['Panes'][0], width=TMUX_SIDE_WIDTH) # Bail early if not (layout['Workers']['Panes'] and layout['Workers']['height']): return # Update worker heights worker_height = layout['Workers']['height'] workers = layout['Workers']['Panes'].copy() num_workers = len(workers) avail_height = sum(tmux.get_pane_size(pane)[1] for pane in workers) avail_height += tmux.get_pane_size()[1] # Current pane # Check if window is too small if avail_height < (worker_height*num_workers) + 3: # Just leave things as-is return # Resize current pane tmux.resize_pane(height=avail_height-(worker_height*num_workers)) # Resize bottom pane tmux.resize_pane(workers.pop(0), height=worker_height) # Resize the rest of the panes by adjusting the ones above them while len(workers) > 1: next_height = sum(tmux.get_pane_size(pane)[1] for pane in workers[:2]) next_height -= worker_height tmux.resize_pane(workers[1], height=next_height) workers.pop(0) def layout_needs_fixed(layout): """Check if layout needs fixed, returns bool.""" needs_fixed = False # Check panes for data in layout.values(): if 'height' in data: needs_fixed = needs_fixed or any( tmux.get_pane_size(pane)[1] != data['height'] for pane in data['Panes'] ) if 'width' in data: needs_fixed = needs_fixed or any( tmux.get_pane_size(pane)[0] != data['width'] for pane in data['Panes'] ) # TODO: Re-enable? ## Group panes #for group in ('Title', 'Info'): # num_panes = len(layout[group]['Panes']) # if num_panes <= 1: # continue # width = int( (tmux.get_pane_size()[0] - (1 - num_panes)) / num_panes ) # for pane in layout[group]['Panes']: # needs_fixed = needs_fixed or abs(tmux.get_pane_size(pane)[0] - width) > 2 # Done return needs_fixed def hmm(): """Hmm?""" def test(): ui = TUI() ui.add_info_pane(height=10, text='Info One') ui.add_info_pane(height=10, text='Info Two') ui.add_info_pane(height=10, text='Info Three') ui.add_worker_pane(height=3, text='Work One') ui.add_worker_pane(height=3, text='Work Two') ui.add_worker_pane(height=3, text='Work Three') return ui if __name__ == '__main__': print("This file is not meant to be called directly.")