373 lines
10 KiB
Python
373 lines
10 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 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': [], 'width': TMUX_SIDE_WIDTH},
|
|
'Progress': {'Panes': [], 'width': TMUX_SIDE_WIDTH},
|
|
}
|
|
|
|
|
|
# Classes
|
|
class TUI():
|
|
"""Object for tracking TUI elements."""
|
|
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'
|
|
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,
|
|
lines: int | None = None,
|
|
percent: int | None = None,
|
|
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 | None = None,
|
|
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,
|
|
'target_id': None,
|
|
'vertical': True,
|
|
})
|
|
|
|
# Update layout
|
|
if update_layout:
|
|
self.layout['Workers']['height'] = height
|
|
|
|
# Add pane
|
|
self.layout['Workers']['Panes'].append(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, 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_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
|
|
|
|
def test() -> TUI:
|
|
"""TODO: Deleteme"""
|
|
ui = TUI()
|
|
ui.add_info_pane(lines=10, text='Info One')
|
|
ui.add_info_pane(lines=10, text='Info Two')
|
|
ui.add_info_pane(lines=10, text='Info Three')
|
|
ui.add_worker_pane(lines=3, text='Work One')
|
|
ui.add_worker_pane(lines=3, text='Work Two')
|
|
ui.add_worker_pane(lines=3, text='Work Three')
|
|
ui.fix_layout()
|
|
return ui
|
|
|
|
|
|
if __name__ == '__main__':
|
|
print("This file is not meant to be called directly.")
|