commands: emit types from partial parser, implement choice completion

mitmproxy/command.py

@@ -15,6 +15,14 @@ from mitmproxy import exceptions
 from mitmproxy import flow
 
 
+def lexer(s):
+    # mypy mis-identifies shlex.shlex as abstract
+    lex = shlex.shlex(s, punctuation_chars=True)  # type: ignore
+    lex.whitespace_split = True
+    lex.commenters = ''
+    return lex
+
+
 Cuts = typing.Sequence[
     typing.Sequence[typing.Union[str, bytes]]
 ]
@@ -123,9 +131,10 @@ class Command:
         return ret
 
 
-class ParseResult(typing.NamedTuple):
+ParseResult = typing.NamedTuple(
-    value: str
+    "ParseResult",
-    type: typing.Type
+    [("value", str), ("type", typing.Type)],
+)
 
 
 class CommandManager:
@@ -147,27 +156,37 @@ class CommandManager:
         """
             Parse a possibly partial command. Return a sequence of (part, type) tuples.
         """
-        parts: typing.List[ParseResult] = []
         buf = io.StringIO(cmdstr)
-        # mypy mis-identifies shlex.shlex as abstract
+        parts: typing.List[str] = []
-        lex = shlex.shlex(buf)  # type: ignore
+        lex = lexer(buf)
         while 1:
             remainder = cmdstr[buf.tell():]
             try:
                 t = lex.get_token()
             except ValueError:
-                parts.append(ParseResult(value = remainder, type = str))
+                parts.append(remainder)
                 break
             if not t:
                 break
-            typ: type = str
+            parts.append(t)
-            # First value is a special case: it has to be a command
         if not parts:
-                typ = Cmd
+            parts = [""]
-            parts.append(ParseResult(value = t, type = typ))
+        elif cmdstr.endswith(" "):
-        if not parts:
+            parts.append("")
-            return [ParseResult(value = "", type = Cmd)]
+
-        return parts
+        parse: typing.List[ParseResult] = []
+        params: typing.List[type] = []
+        for i in range(len(parts)):
+            if i == 0:
+                params[:] = [Cmd]
+                if parts[i] in self.commands:
+                    params.extend(self.commands[parts[i]].paramtypes)
+            if params:
+                typ = params.pop(0)
+            else:
+                typ = str
+            parse.append(ParseResult(value=parts[i], type=typ))
+        return parse
 
     def call_args(self, path, args):
         """
@@ -181,7 +200,7 @@ class CommandManager:
         """
             Call a command using a string. May raise CommandError.
         """
-        parts = shlex.split(cmdstr)
+        parts = list(lexer(cmdstr))
         if not len(parts) >= 1:
             raise exceptions.CommandError("Invalid command: %s" % cmdstr)
         return self.call_args(parts[0], parts[1:])
@@ -208,8 +227,6 @@ def parsearg(manager: CommandManager, spec: str, argtype: type) -> typing.Any:
                 "Invalid choice: see %s for options" % cmd
             )
         return spec
-    if argtype in (Path, Cmd):
-        return spec
     elif issubclass(argtype, str):
         return spec
     elif argtype == bool:

mitmproxy/tools/console/commander/commander.py

@@ -72,7 +72,8 @@ class CommandBuffer():
     def cycle_completion(self) -> None:
         if not self.completion:
             parts = self.master.commands.parse_partial(self.buf[:self.cursor])
-            if parts[-1].type == mitmproxy.command.Cmd:
+            last = parts[-1]
+            if last.type == mitmproxy.command.Cmd:
                 self.completion = CompletionState(
                     completer = ListCompleter(
                         parts[-1].value,
@@ -80,6 +81,14 @@ class CommandBuffer():
                     ),
                     parse = parts,
                 )
+            elif isinstance(last.type, mitmproxy.command.Choice):
+                self.completion = CompletionState(
+                    completer = ListCompleter(
+                        parts[-1].value,
+                        self.master.commands.call(last.type.options_command),
+                    ),
+                    parse = parts,
+                )
         if self.completion:
             nxt = self.completion.completer.cycle()
             buf = " ".join([i.value for i in self.completion.parse[:-1]]) + " " + nxt

test/mitmproxy/test_command.py

@@ -11,19 +11,24 @@ from mitmproxy.utils import typecheck
 
 
 class TAddon:
+    @command.command("cmd1")
     def cmd1(self, foo: str) -> str:
         """cmd1 help"""
         return "ret " + foo
 
+    @command.command("cmd2")
     def cmd2(self, foo: str) -> str:
         return 99
 
+    @command.command("cmd3")
     def cmd3(self, foo: int) -> int:
         return foo
 
+    @command.command("empty")
     def empty(self) -> None:
         pass
 
+    @command.command("varargs")
     def varargs(self, one: str, *var: str) -> typing.Sequence[str]:
         return list(var)
 
@@ -34,6 +39,7 @@ class TAddon:
     def choose(self, arg: str) -> typing.Sequence[str]:
         return ["one", "two", "three"]
 
+    @command.command("path")
     def path(self, arg: command.Path) -> None:
         pass
 
@@ -82,11 +88,25 @@ class TestCommand:
             ],
             ["a", [command.ParseResult(value = "a", type = command.Cmd)]],
             ["", [command.ParseResult(value = "", type = command.Cmd)]],
+            [
+                "cmd3 1",
+                [
+                    command.ParseResult(value = "cmd3", type = command.Cmd),
+                    command.ParseResult(value = "1", type = int),
+                ]
+            ],
+            [
+                "cmd3 ",
+                [
+                    command.ParseResult(value = "cmd3", type = command.Cmd),
+                    command.ParseResult(value = "", type = int),
+                ]
+            ],
         ]
         with taddons.context() as tctx:
-            cm = command.CommandManager(tctx.master)
+            tctx.master.addons.add(TAddon())
             for s, expected in tests:
-                assert cm.parse_partial(s) == expected
+                assert tctx.master.commands.parse_partial(s) == expected
 
 
 def test_simple():