From 196e2adc824459e828f7f09a5eb59ad388b44c47 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sun, 10 Nov 2019 19:10:35 -0700
Subject: [PATCH] Added tmux layout maintenance sections

* Support both threading and signal based calls
  * Should provide a smoother UIX under Linux & macOS
---
 scripts/wk/cfg/hw.py   | 13 ++++++++++
 scripts/wk/hw/diags.py | 32 +++++++++++++++++++++++-
 scripts/wk/tmux.py     | 56 ++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 98 insertions(+), 3 deletions(-)

diff --git a/scripts/wk/cfg/hw.py b/scripts/wk/cfg/hw.py
index 2247b8dc..1e848eb6 100644
--- a/scripts/wk/cfg/hw.py
+++ b/scripts/wk/cfg/hw.py
@@ -4,6 +4,8 @@
 
 import re
 
+from collections import OrderedDict
+
 
 # STATIC VARIABLES
 ATTRIBUTE_COLORS = (
@@ -53,6 +55,17 @@ REGEX_POWER_ON_TIME = re.compile(
   r'^(\d+)([Hh].*|\s+\(\d+\s+\d+\s+\d+\).*)'
   )
 TMUX_SIDE_WIDTH = 20
+TMUX_LAYOUT = OrderedDict({
+  'Top':            {'height':  2,                'Check': True},
+  'Started':        {'width':   TMUX_SIDE_WIDTH,  'Check': True},
+  'Progress':       {'width':   TMUX_SIDE_WIDTH,  'Check': True},
+  # Testing panes
+  'Prime95':        {'height':  11,               'Check': False},
+  'Temps':          {'height':  1000,             'Check': False},
+  'SMART':          {'height':  3,                'Check': True},
+  'badblocks':      {'height':  5,                'Check': True},
+  'I/O Benchmark':  {'height':  1000,             'Check': False},
+  })
 
 
 if __name__ == '__main__':
diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py
index 6405cf05..8e8a301f 100644
--- a/scripts/wk/hw/diags.py
+++ b/scripts/wk/hw/diags.py
@@ -6,13 +6,14 @@ import logging
 import os
 import pathlib
 import platform
+import signal
 import time
 
 from collections import OrderedDict
 from docopt import docopt
 
 from wk import exe, net, std, tmux
-from wk.cfg.hw import TMUX_SIDE_WIDTH
+from wk.cfg.hw import TMUX_LAYOUT, TMUX_SIDE_WIDTH
 from wk.cfg.main import KIT_NAME_FULL
 
 
@@ -100,7 +101,36 @@ class State():
         },
       })
     self.top_text = std.color_string('Hardware Diagnostics', 'GREEN')
+
+    # Init tmux and start a background process to maintain layout
     self.init_tmux()
+    if hasattr(signal, 'SIGWINCH'):
+      # Use signal handling
+      signal.signal(signal.SIGWINCH, self.fix_tmux_layout)
+    else:
+      exe.start_thread(self.fix_tmux_layout_loop)
+
+  def fix_tmux_layout(self, forced=True, signum=None, frame=None):
+    # pylint: disable=unused-argument
+    """Fix tmux layout based on TMUX_LAYOUT.
+
+    NOTE: To support being called by both a signal and a thread
+          signum and frame must be valid aguments.
+    """
+    try:
+      tmux.fix_layout(self.panes, TMUX_LAYOUT, forced=forced)
+    except RuntimeError:
+      # Assuming self.panes changed while running
+      pass
+
+  def fix_tmux_layout_loop(self):
+    """Fix tmux layout on a loop.
+
+    NOTE: This should be called as a thread.
+    """
+    while True:
+      self.fix_tmux_layout(forced=False)
+      std.sleep(1)
 
   def init_tmux(self):
     """Initialize tmux layout."""
diff --git a/scripts/wk/tmux.py b/scripts/wk/tmux.py
index 2b726c2a..84068ab5 100644
--- a/scripts/wk/tmux.py
+++ b/scripts/wk/tmux.py
@@ -24,6 +24,27 @@ def capture_pane(pane_id=None):
   return proc.stdout.strip()
 
 
+def fix_layout(panes, layout, forced=False):
+  """Fix pane sizes based on layout."""
+  if not (forced or layout_needs_fixed(panes, layout)):
+    # Layout should be fine
+    return
+
+  # Update panes
+  for name, data in layout.items():
+    # Skip missing panes
+    if name not in panes:
+      continue
+
+    # Resize pane
+    pane_id = panes[name]
+    try:
+      resize_pane(pane_id, **data)
+    except RuntimeError:
+      # Assuming pane was closed just before resizing
+      pass
+
+
 def get_pane_size(pane_id=None):
   """Get current or target pane size, returns tuple."""
   cmd = ['tmux', 'display', '-p']
@@ -60,6 +81,32 @@ def kill_pane(*pane_ids):
     run_program(cmd+[pane_id], check=False)
 
 
+def layout_needs_fixed(panes, layout):
+  """Check if layout needs fixed, returns bool."""
+  needs_fixed = False
+
+  # Check panes
+  for name, data in layout.items():
+    # Skip unpredictably sized panes
+    if not data.get('Check', False):
+      continue
+
+    # Skip missing panes
+    if name not in panes:
+      continue
+
+    # Check pane size
+    pane_id = panes[name]
+    width, height = get_pane_size(pane_id)
+    if data.get('width', False) and data['width'] != width:
+      needs_fixed = True
+    if data.get('height', False) and data['height'] != height:
+      needs_fixed = True
+
+  # Done
+  return needs_fixed
+
+
 def poll_pane(pane_id):
   """Check if pane exists, returns bool."""
   cmd = ['tmux', 'list-panes', '-F', '#D']
@@ -133,8 +180,13 @@ def prep_file(path):
     pass
 
 
-def resize_pane(pane_id=None, width=None, height=None):
-  """Resize current or target pane."""
+def resize_pane(pane_id=None, width=None, height=None, **kwargs):
+  # pylint: disable=unused-argument
+  """Resize current or target pane.
+
+  NOTE: kwargs is only here to make calling this function easier
+        by dropping any extra kwargs passed.
+  """
   cmd = ['tmux', 'resize-pane']
 
   # Safety checks