console options: handle choices settings

This implements a generic chooser overlay, and uses it to handle setting
options that have fixed choices. We'll use this overlay elsewhere too.

3 files changed, 122 insertions, 1 deletions

diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index d0e23712..5d481eeb 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -24,6 +24,7 @@ from mitmproxy.tools.console import flowview
 from mitmproxy.tools.console import grideditor
 from mitmproxy.tools.console import help
 from mitmproxy.tools.console import options
+from mitmproxy.tools.console import overlay
 from mitmproxy.tools.console import palettepicker
 from mitmproxy.tools.console import palettes
 from mitmproxy.tools.console import signals
@@ -285,7 +286,6 @@ class ConsoleMaster(master.Master):
         self.ab = statusbar.ActionBar()
 
         self.loop.set_alarm_in(0.01, self.ticker)
-
         self.loop.set_alarm_in(
             0.0001,
             lambda *args: self.view_flowlist()
@@ -309,6 +309,16 @@ class ConsoleMaster(master.Master):
     def shutdown(self):
         raise urwid.ExitMainLoop
 
+    def overlay(self, widget):
+        signals.push_view_state.send(
+            self,
+            window = overlay.SimpleOverlay(
+                widget,
+                self.loop.widget,
+                widget.width,
+            )
+        )
+
     def view_help(self, helpctx):
         signals.push_view_state.send(
             self,
diff --git a/mitmproxy/tools/console/options.py b/mitmproxy/tools/console/options.py
index 5458bc76..479d3918 100644
--- a/mitmproxy/tools/console/options.py
+++ b/mitmproxy/tools/console/options.py
@@ -6,6 +6,7 @@ from typing import Optional
 from mitmproxy import exceptions
 from mitmproxy.tools.console import common
 from mitmproxy.tools.console import signals
+from mitmproxy.tools.console import overlay
 
 
 def can_edit_inplace(opt):
@@ -121,6 +122,7 @@ class OptionListWalker(urwid.ListWalker):
 
     def sig_mod(self, *args, **kwargs):
         self._modified()
+        self.set_focus(self.index)
 
     def start_editing(self):
         self.editing = True
@@ -202,6 +204,15 @@ class OptionsList(urwid.ListBox):
                 elif can_edit_inplace(foc.opt):
                     self.walker.start_editing()
                     self.walker._modified()
+                elif foc.opt.choices:
+                    self.master.overlay(
+                        overlay.Chooser(
+                            foc.opt.name,
+                            foc.opt.choices,
+                            foc.opt.current(),
+                            self.master.options.setter(foc.opt.name)
+                        )
+                    )
         return super().keypress(size, key)
 
 
diff --git a/mitmproxy/tools/console/overlay.py b/mitmproxy/tools/console/overlay.py
new file mode 100644
index 00000000..cf7fdfcb
--- /dev/null
+++ b/mitmproxy/tools/console/overlay.py
@@ -0,0 +1,100 @@
+from mitmproxy.tools.console import common
+from mitmproxy.tools.console import signals
+import urwid
+
+
+class SimpleOverlay(urwid.Overlay):
+    def __init__(self, widget, parent, width):
+        super().__init__(
+            widget,
+            parent,
+            align="center",
+            width=width,
+            valign="middle",
+            height="pack"
+        )
+
+    def keypress(self, size, key):
+        if key == "esc":
+            signals.pop_view_state.send(self)
+        return super().keypress(size, key)
+
+
+class Choice(urwid.WidgetWrap):
+    def __init__(self, txt, focus, current):
+        if current:
+            s = "option_active_selected" if focus else "option_active"
+        else:
+            s = "option_selected" if focus else "text"
+        return super().__init__(
+            urwid.AttrWrap(
+                urwid.Padding(urwid.Text(txt)),
+                s,
+            )
+        )
+
+    def selectable(self):
+        return True
+
+    def keypress(self, size, key):
+        return key
+
+
+class ChooserListWalker(urwid.ListWalker):
+    def __init__(self, choices, current):
+        self.index = 0
+        self.choices = choices
+        self.current = current
+
+    def _get(self, idx, focus):
+        c = self.choices[idx]
+        return Choice(c, focus, c == self.current)
+
+    def set_focus(self, index):
+        self.index = index
+
+    def get_focus(self):
+        return self._get(self.index, True), self.index
+
+    def get_next(self, pos):
+        if pos >= len(self.choices) - 1:
+            return None, None
+        pos = pos + 1
+        return self._get(pos, False), pos
+
+    def get_prev(self, pos):
+        pos = pos - 1
+        if pos < 0:
+            return None, None
+        return self._get(pos, False), pos
+
+
+class Chooser(urwid.WidgetWrap):
+    def __init__(self, title, choices, current, callback):
+        self.choices = choices
+        self.callback = callback
+        choicewidth = max([len(i) for i in choices])
+        self.width = max(choicewidth, len(title) + 5)
+        self.walker = ChooserListWalker(choices, current)
+        super().__init__(
+            urwid.AttrWrap(
+                urwid.LineBox(
+                    urwid.BoxAdapter(
+                        urwid.ListBox(self.walker),
+                        len(choices)
+                    ),
+                    title= title
+                ),
+                "background"
+            )
+        )
+
+    def selectable(self):
+        return True
+
+    def keypress(self, size, key):
+        key = common.shortcuts(key)
+        if key == "enter":
+            self.callback(self.choices[self.walker.index])
+            signals.pop_view_state.send(self)
+        return super().keypress(size, key)
\ No newline at end of file