From 158c3f76b66f305cb3a1f315e929e4cc457d7b8d Mon Sep 17 00:00:00 2001
From: Jack Foltz <jack@foltz.io>
Date: Thu, 6 Dec 2018 01:30:41 -0500
Subject: [PATCH] Add CheckList menu and menu titles

---
 lib/menu.py | 89 ++++++++++++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 71 insertions(+), 18 deletions(-)

diff --git a/lib/menu.py b/lib/menu.py
index 6771aeb..6365f03 100755
--- a/lib/menu.py
+++ b/lib/menu.py
@@ -1,5 +1,6 @@
 import curses
 from curses import panel
+from types import LambdaType
 
 class BaseMenu:
     def __init__(self, stdscreen):
@@ -25,12 +26,13 @@ class BaseMenu:
         curses.doupdate()
 
 class Menu(BaseMenu):
-    def __init__(self, stdscreen, items):
+    def __init__(self, stdscreen, title, items):
         super().__init__(stdscreen)
         self.should_exit = False
         self.position = 0
+        self.title = title
         self.items = items
-        self.items.append(('exit', self.exit))
+        self.items.append(('Exit', self.exit))
 
     def exit(self):
         self.should_exit = True
@@ -39,30 +41,44 @@ class Menu(BaseMenu):
         self.position += offset
         self.position %= len(self.items)
 
+    def draw_title(self):
+        self.window.addstr(1, 1, self.title, curses.A_BOLD)
+
     def draw(self):
+        self.draw_title()
         for index, item in enumerate(self.items):
             text = '%d. %s' % (index + 1, item[0])
             text_mode = curses.A_REVERSE if index == self.position else curses.A_NORMAL
-            self.window.addstr(index + 1, 1, text, text_mode)
+            self.window.addstr(index + 3, 1, text, text_mode)
 
-    def get_selection(self, key):
+    def select(self, index):
+        self.items[index][1]()
+            
+    def input(self, key):
+        # ESC to quit
+        if key == 27:
+            self.window.nodelay(True)
+            key = self.window.getch()
+            if key == -1:
+                return self.exit()
+
+        # q to quit
+        if key == ord('q'):
+            return self.exit()
+
+        # Enter to select item
         if key in [curses.KEY_ENTER, ord('\n')]:
-            return self.items[self.position]
-        elif (key - ord('0')) in range(1, len(self.items) + 1):
-            return self.items[key - ord('0') - 1]
-        else:
-            return None
+            self.select(self.position)
 
-    def input(self):
-        key = self.window.getch()
-        selection = self.get_selection(key)
+        # Number keys to select item
+        if key - ord('0') in range(1, len(self.items) + 1):
+            self.select(key - ord('0') - 1)
 
-        if selection:
-            selection[1]()
-
-        if key == curses.KEY_UP:
+        # Arrow keys or jk to navigate
+        if key in [curses.KEY_UP, ord('k')]:
             self.navigate(-1)
-        elif key == curses.KEY_DOWN:
+
+        elif key in [curses.KEY_DOWN, ord('j')]:
             self.navigate(1)
 
     def display(self):
@@ -70,7 +86,44 @@ class Menu(BaseMenu):
         while not self.should_exit:
             self.update()
             self.draw()
-            self.input()
+            key = self.window.getch()
+            self.input(key)
 
         self.should_exit = False
         self.down()
+
+class ChecklistMenu(Menu):
+    def __init__(self, stdscreen, title, items, default):
+        super().__init__(stdscreen, title, items)
+        self.states = [default] * len(items)
+    
+    def draw(self):
+        self.draw_title()
+        for index, item in enumerate(self.items):
+            text = ''
+            if isinstance(item[1], LambdaType):
+                text = '%d. [ %s ] %s' % (index + 1, '*' if self.states[index] else ' ', item[0])
+            else:
+                text = '%d. %s' % (index + 1, item[0])
+
+            text_mode = curses.A_REVERSE if index == self.position else curses.A_NORMAL
+            self.window.addstr(index + 3, 1, text, text_mode)
+
+    def input(self, key):
+        super().input(key)
+
+        if key == ord('\t'):
+            self.expand(self.position)
+
+    def expand(self, index):
+        if len(self.items[index]) >= 3:
+            self.items[index][2]()
+
+    def select(self, index):
+        self.states[index] = not self.states[index]
+        fn = self.items[index][1]
+        if isinstance(fn, LambdaType):
+            fn(self.states[index])
+        else:
+            fn()
+