WizardKit/scripts/wk/ui/tui.py
2023-04-09 16:32:13 -07:00

304 lines
8.6 KiB
Python

"""WizardKit: TUI functions"""
# vim: sts=2 sw=2 ts=2
import atexit
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)
# Close all panes at exit
atexit.register(tmux.kill_all_panes)
def add_info_pane(self, height, **tmux_args) -> None:
"""Add info pane."""
tmux_args.update({
'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.")