diff options
-rw-r--r-- | mitmproxy/command.py | 9 | ||||
-rw-r--r-- | test/mitmproxy/test_command.py | 19 |
2 files changed, 26 insertions, 2 deletions
diff --git a/mitmproxy/command.py b/mitmproxy/command.py index 087f7770..05caf261 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -181,13 +181,18 @@ class CommandManager: parse = [] # type: typing.List[ParseResult] params = [] # type: typing.List[type] + typ = None # type: typing.Type for i in range(len(parts)): if i == 0: - params[:] = [Cmd] + typ = Cmd if parts[i] in self.commands: params.extend(self.commands[parts[i]].paramtypes) - if params: + elif params: typ = params.pop(0) + # FIXME: Do we need to check that Arg is positional? + if typ == Cmd and params and params[0] == Arg: + if parts[i] in self.commands: + params[:] = self.commands[parts[i]].paramtypes else: typ = str parse.append(ParseResult(value=parts[i], type=typ)) diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py index 76ce2245..298b34fb 100644 --- a/test/mitmproxy/test_command.py +++ b/test/mitmproxy/test_command.py @@ -24,6 +24,10 @@ class TAddon: def cmd3(self, foo: int) -> int: return foo + @command.command("subcommand") + def subcommand(self, cmd: command.Cmd, *args: command.Arg) -> str: + return "ok" + @command.command("empty") def empty(self) -> None: pass @@ -102,6 +106,21 @@ class TestCommand: command.ParseResult(value = "", type = int), ] ], + [ + "subcommand ", + [ + command.ParseResult(value = "subcommand", type = command.Cmd), + command.ParseResult(value = "", type = command.Cmd), + ] + ], + [ + "subcommand cmd3 ", + [ + command.ParseResult(value = "subcommand", type = command.Cmd), + command.ParseResult(value = "cmd3", type = command.Cmd), + command.ParseResult(value = "", type = int), + ] + ], ] with taddons.context() as tctx: tctx.master.addons.add(TAddon()) |