commander: tab completion for command names

3 files changed, 77 insertions, 10 deletions

diff --git a/mitmproxy/command.py b/mitmproxy/command.py
index aab721b5..b909dfd5 100644
--- a/mitmproxy/command.py
+++ b/mitmproxy/command.py
@@ -125,7 +125,7 @@ class Command:
 
 class ParseResult(typing.NamedTuple):
     value: str
-    type: type
+    type: typing.Type
 
 
 class CommandManager:
diff --git a/mitmproxy/tools/console/commander/commander.py b/mitmproxy/tools/console/commander/commander.py
index d961d421..a0c3a3b2 100644
--- a/mitmproxy/tools/console/commander/commander.py
+++ b/mitmproxy/tools/console/commander/commander.py
@@ -1,23 +1,51 @@
 import urwid
 from urwid.text_layout import calc_coords
 import typing
+import abc
 
 import mitmproxy.master
 import mitmproxy.command
 
 
-class CompletionState:
-    def __init__(self, parts: typing.Sequence[mitmproxy.command.ParseResult]) -> None:
-        self.parts = parts
+class Completer:
+    @abc.abstractmethod
+    def cycle(self) -> str:
+        pass
+
+
+class ListCompleter(Completer):
+    def __init__(
+        self,
+        start: str,
+        options: typing.Sequence[str],
+    ) -> None:
+        self.start = start
+        self.options = []  # type: typing.Sequence[str]
+        for o in options:
+            if o.startswith(start):
+                self.options.append(o)
+        self.offset = 0
+
+    def cycle(self) -> str:
+        if not self.options:
+            return self.start
+        ret = self.options[self.offset]
+        self.offset = (self.offset + 1) % len(self.options)
+        return ret
+
+
+class CompletionState(typing.NamedTuple):
+    completer: Completer
+    parse: typing.Sequence[mitmproxy.command.ParseResult]
 
 
 class CommandBuffer():
     def __init__(self, master: mitmproxy.master.Master, start: str = "") -> None:
         self.master = master
         self.buf = start
-        # This is the logical cursor position - the display cursor is one
-        # character further on. Cursor is always within the range [0:len(buffer)].
+        # Cursor is always within the range [0:len(buffer)].
         self._cursor = len(self.buf)
+        self.completion = None   # type: CompletionState
 
     @property
     def cursor(self) -> int:
@@ -42,16 +70,29 @@ class CommandBuffer():
         self.cursor = self.cursor + 1
 
     def cycle_completion(self) -> None:
-        parts = self.master.commands.parse_partial(self.buf[:self.cursor])
-        if parts[-1][1] == str:
-            return
-        raise ValueError
+        if not self.completion:
+            parts = self.master.commands.parse_partial(self.buf[:self.cursor])
+            if parts[-1].type == mitmproxy.command.Cmd:
+                self.completion = CompletionState(
+                    completer = ListCompleter(
+                        parts[-1].value,
+                        self.master.commands.commands.keys(),
+                    ),
+                    parse = parts,
+                )
+        if self.completion:
+            nxt = self.completion.completer.cycle()
+            buf = " ".join([i.value for i in self.completion.parse[:-1]]) + " " + nxt
+            buf = buf.strip()
+            self.buf = buf
+            self.cursor = len(self.buf)
 
     def backspace(self) -> None:
         if self.cursor == 0:
             return
         self.buf = self.buf[:self.cursor - 1] + self.buf[self.cursor:]
         self.cursor = self.cursor - 1
+        self.completion = None
 
     def insert(self, k: str) -> None:
         """
@@ -59,6 +100,7 @@ class CommandBuffer():
         """
         self.buf = self.buf = self.buf[:self.cursor] + k + self.buf[self.cursor:]
         self.cursor += 1
+        self.completion = None
 
 
 class CommandEdit(urwid.WidgetWrap):
diff --git a/test/mitmproxy/tools/console/test_commander.py b/test/mitmproxy/tools/console/test_commander.py
index 9ef4a318..1ac4c5c6 100644
--- a/test/mitmproxy/tools/console/test_commander.py
+++ b/test/mitmproxy/tools/console/test_commander.py
@@ -2,6 +2,31 @@ from mitmproxy.tools.console.commander import commander
 from mitmproxy.test import taddons
 
 
+class TestListCompleter:
+    def test_cycle(self):
+        tests = [
+            [
+                "",
+                ["a", "b", "c"],
+                ["a", "b", "c", "a"]
+            ],
+            [
+                "xxx",
+                ["a", "b", "c"],
+                ["xxx", "xxx", "xxx"]
+            ],
+            [
+                "b",
+                ["a", "b", "ba", "bb", "c"],
+                ["b", "ba", "bb", "b"]
+            ],
+        ]
+        for start, options, cycle in tests:
+            c = commander.ListCompleter(start, options)
+            for expected in cycle:
+                assert c.cycle() == expected
+
+
 class TestCommandBuffer:
 
     def test_backspace(self):