2Shirt
4a34f5477d
Avoids issue where the main menu is printed before the layout is fully set causing the first few lines to be hidden by the title pane.
382 lines
11 KiB
Python
382 lines
11 KiB
Python
"""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 typing import Any
|
|
|
|
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': [], 'height': TMUX_TITLE_HEIGHT},
|
|
'Progress': {'Panes': [], 'width': TMUX_SIDE_WIDTH},
|
|
}
|
|
|
|
|
|
# Classes
|
|
class TUI():
|
|
"""Object for tracking TUI elements."""
|
|
def __init__(self, title_text: str | None = None):
|
|
self.clear_on_resize = False
|
|
self.layout: dict[str, dict[str, Any]] = deepcopy(TMUX_LAYOUT)
|
|
self.side_width: int = TMUX_SIDE_WIDTH
|
|
self.title_text: str = title_text if title_text else 'Title Text'
|
|
self.title_text_line2: str = ''
|
|
self.title_colors: list[str] = ['BLUE', '']
|
|
|
|
# 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: int | None = None,
|
|
percent: int = 0,
|
|
update_layout: bool = 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: 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()
|
|
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: int | None = None,
|
|
percent: int = 0,
|
|
update_layout: bool = 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 if percent else None,
|
|
'target_id': None,
|
|
'vertical': True,
|
|
})
|
|
|
|
# Update layout
|
|
if update_layout:
|
|
self.layout['Workers']['height'] = height
|
|
|
|
# Add pane (ensure panes are sorted top to bottom)
|
|
self.layout['Workers']['Panes'].insert(0, 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: bool = True) -> None:
|
|
"""Fix tmux layout based on self.layout."""
|
|
try:
|
|
tmux.fix_layout(self.layout, clear_on_resize=self.clear_on_resize, 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))
|
|
|
|
# Progress
|
|
self.layout['Progress']['Panes'].append(tmux.split_window(
|
|
lines=TMUX_SIDE_WIDTH,
|
|
text=' ',
|
|
))
|
|
|
|
# Started
|
|
self.layout['Started']['Panes'].append(tmux.split_window(
|
|
behind=True,
|
|
lines=2,
|
|
target_id=self.layout['Progress']['Panes'][0],
|
|
text=ansi.color_string(
|
|
['Started', time.strftime("%Y-%m-%d %H:%M %Z")],
|
|
['BLUE', None],
|
|
sep='\n',
|
|
),
|
|
vertical=True,
|
|
))
|
|
|
|
# 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',
|
|
),
|
|
))
|
|
|
|
# Done
|
|
sleep(0.2)
|
|
|
|
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 reset_title_pane(
|
|
self,
|
|
line1: str = 'Title Text',
|
|
line2: str = '',
|
|
colors: list[str] | None = None,
|
|
) -> None:
|
|
"""Remove all extra title panes, reset main title pane, and update layout."""
|
|
colors = self.title_colors if colors is None else colors
|
|
panes = self.layout['Title']['Panes'].copy()
|
|
if len(panes) > 1:
|
|
tmux.kill_pane(*panes[1:])
|
|
self.layout['Title']['Panes'] = panes[:1]
|
|
self.set_title(line1, line2, colors)
|
|
|
|
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: 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: 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 ''
|
|
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: bool = False) -> None:
|
|
"""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) -> bool:
|
|
"""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
|
|
|
|
|
|
if __name__ == '__main__':
|
|
print("This file is not meant to be called directly.")
|