WIP: autocompletion

mitmproxy/command.py

@@ -123,6 +123,11 @@ class Command:
         return ret
 
 
+class ParseResult(typing.NamedTuple):
+    value: str
+    type: type
+
+
 class CommandManager:
     def __init__(self, master):
         self.master = master
@@ -138,11 +143,11 @@ class CommandManager:
     def add(self, path: str, func: typing.Callable):
         self.commands[path] = Command(self, path, func)
 
-    def parse_partial(self, cmdstr: str) -> typing.Sequence[typing.Tuple[str, type]]:
+    def parse_partial(self, cmdstr: str) -> typing.Sequence[ParseResult]:
         """
             Parse a possibly partial command. Return a sequence of (part, type) tuples.
         """
-        parts: typing.List[typing.Tuple[str, type]] = []
+        parts: typing.List[ParseResult] = []
         buf = io.StringIO(cmdstr)
         # mypy mis-identifies shlex.shlex as abstract
         lex = shlex.shlex(buf)  # type: ignore
@@ -151,7 +156,7 @@ class CommandManager:
             try:
                 t = lex.get_token()
             except ValueError:
-                parts.append((remainder, str))
+                parts.append(ParseResult(value = remainder, type = str))
                 break
             if not t:
                 break
@@ -159,7 +164,9 @@ class CommandManager:
             # First value is a special case: it has to be a command
             if not parts:
                 typ = Cmd
-            parts.append((t, typ))
+            parts.append(ParseResult(value = t, type = typ))
+        if not parts:
+            return [ParseResult(value = "", type = Cmd)]
         return parts
 
     def call_args(self, path, args):

mitmproxy/tools/console/commander/commander.py

@@ -1,9 +1,19 @@
 import urwid
 from urwid.text_layout import calc_coords
+import typing
+
+import mitmproxy.master
+import mitmproxy.command
+
+
+class CompletionState:
+    def __init__(self, parts: typing.Sequence[mitmproxy.command.ParseResult]) -> None:
+        self.parts = parts
 
 
 class CommandBuffer():
-    def __init__(self, start: str = "") -> None:
+    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)].
@@ -31,6 +41,12 @@ class CommandBuffer():
     def right(self) -> None:
         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
+
     def backspace(self) -> None:
         if self.cursor == 0:
             return
@@ -48,8 +64,9 @@ class CommandBuffer():
 class CommandEdit(urwid.WidgetWrap):
     leader = ": "
 
-    def __init__(self, text) -> None:
-        self.cbuf = CommandBuffer(text)
+    def __init__(self, master: mitmproxy.master.Master, text: str) -> None:
+        self.master = master
+        self.cbuf = CommandBuffer(master, text)
         self._w = urwid.Text(self.leader)
         self.update()
 
@@ -60,6 +77,8 @@ class CommandEdit(urwid.WidgetWrap):
             self.cbuf.left()
         elif key == "right":
             self.cbuf.right()
+        elif key == "tab":
+            self.cbuf.cycle_completion()
         elif len(key) == 1:
             self.cbuf.insert(key)
         self.update()

mitmproxy/tools/console/statusbar.py

@@ -67,7 +67,7 @@ class ActionBar(urwid.WidgetWrap):
 
     def sig_prompt_command(self, sender, partial=""):
         signals.focus.send(self, section="footer")
-        self._w = commander.CommandEdit(partial)
+        self._w = commander.CommandEdit(self.master, partial)
         self.prompting = commandeditor.CommandExecutor(self.master)
 
     def sig_prompt_onekey(self, sender, prompt, keys, callback, args=()):

test/mitmproxy/test_command.py

@@ -66,8 +66,22 @@ class TestCommand:
 
     def test_parse_partial(self):
         tests = [
-            ["foo bar", [("foo", command.Cmd), ("bar", str)]],
-            ["foo 'bar", [("foo", command.Cmd), ("'bar", str)]],
+            [
+                "foo bar",
+                [
+                    command.ParseResult(value = "foo", type = command.Cmd),
+                    command.ParseResult(value = "bar", type = str)
+                ],
+            ],
+            [
+                "foo 'bar",
+                [
+                    command.ParseResult(value = "foo", type = command.Cmd),
+                    command.ParseResult(value = "'bar", type = str)
+                ]
+            ],
+            ["a", [command.ParseResult(value = "a", type = command.Cmd)]],
+            ["", [command.ParseResult(value = "", type = command.Cmd)]],
         ]
         with taddons.context() as tctx:
             cm = command.CommandManager(tctx.master)

test/mitmproxy/tools/console/test_commander.py

@@ -1,5 +1,5 @@
-
 from mitmproxy.tools.console.commander import commander
+from mitmproxy.test import taddons
 
 
 class TestCommandBuffer:
@@ -13,8 +13,9 @@ class TestCommandBuffer:
             [("123", 2), ("13", 1)],
             [("123", 0), ("123", 0)],
         ]
+        with taddons.context() as tctx:
             for start, output in tests:
-            cb = commander.CommandBuffer()
+                cb = commander.CommandBuffer(tctx.master)
                 cb.buf, cb.cursor = start[0], start[1]
                 cb.backspace()
                 assert cb.buf == output[0]
@@ -26,9 +27,17 @@ class TestCommandBuffer:
             [("a", 0), ("xa", 1)],
             [("xa", 2), ("xax", 3)],
         ]
+        with taddons.context() as tctx:
             for start, output in tests:
-            cb = commander.CommandBuffer()
+                cb = commander.CommandBuffer(tctx.master)
                 cb.buf, cb.cursor = start[0], start[1]
                 cb.insert("x")
                 assert cb.buf == output[0]
                 assert cb.cursor == output[1]
+
+    def test_cycle_completion(self):
+        with taddons.context() as tctx:
+            cb = commander.CommandBuffer(tctx.master)
+            cb.buf = "foo bar"
+            cb.cursor = len(cb.buf)
+            cb.cycle_completion()