commands: support *args for commands

Use this to simplify meta-commands in console, and to create a console_choose
command that prompts the user for a choice, and then executes a command with
variable substitution.

2 files changed, 48 insertions, 28 deletions

diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py
index ac082153..958328b2 100644
--- a/test/mitmproxy/test_command.py
+++ b/test/mitmproxy/test_command.py
@@ -1,9 +1,6 @@
 import typing
 from mitmproxy import command
 from mitmproxy import flow
-from mitmproxy import master
-from mitmproxy import options
-from mitmproxy import proxy
 from mitmproxy import exceptions
 from mitmproxy.test import tflow
 from mitmproxy.test import taddons
@@ -19,24 +16,41 @@ class TAddon:
     def cmd2(self, foo: str) -> str:
         return 99
 
+    def cmd3(self, foo: int) -> int:
+        return foo
+
     def empty(self) -> None:
         pass
 
+    def varargs(self, one: str, *var: typing.Sequence[str]) -> typing.Sequence[str]:
+        return list(var)
+
 
 class TestCommand:
+    def test_varargs(self):
+        with taddons.context() as tctx:
+            cm = command.CommandManager(tctx.master)
+            a = TAddon()
+            c = command.Command(cm, "varargs", a.varargs)
+            assert c.signature_help() == "varargs str *str -> [str]"
+            assert c.call(["one", "two", "three"]) == ["two", "three"]
+            with pytest.raises(exceptions.CommandError):
+                c.call(["one", "two", 3])
+
     def test_call(self):
-        o = options.Options()
-        m = master.Master(o, proxy.DummyServer(o))
-        cm = command.CommandManager(m)
+        with taddons.context() as tctx:
+            cm = command.CommandManager(tctx.master)
+            a = TAddon()
+            c = command.Command(cm, "cmd.path", a.cmd1)
+            assert c.call(["foo"]) == "ret foo"
+            assert c.signature_help() == "cmd.path str -> str"
 
-        a = TAddon()
-        c = command.Command(cm, "cmd.path", a.cmd1)
-        assert c.call(["foo"]) == "ret foo"
-        assert c.signature_help() == "cmd.path str -> str"
+            c = command.Command(cm, "cmd.two", a.cmd2)
+            with pytest.raises(exceptions.CommandError):
+                c.call(["foo"])
 
-        c = command.Command(cm, "cmd.two", a.cmd2)
-        with pytest.raises(exceptions.CommandError):
-            c.call(["foo"])
+            c = command.Command(cm, "cmd.three", a.cmd3)
+            assert c.call(["1"]) == 1
 
 
 def test_simple():
@@ -74,14 +88,12 @@ def test_typename():
 
 
 class DummyConsole:
-    def load(self, l):
-        l.add_command("view.resolve", self.resolve)
-        l.add_command("cut", self.cut)
-
+    @command.command("view.resolve")
     def resolve(self, spec: str) -> typing.Sequence[flow.Flow]:
         n = int(spec)
         return [tflow.tflow(resp=True)] * n
 
+    @command.command("cut")
     def cut(self, spec: str) -> command.Cuts:
         return [["test"]]
 
@@ -115,6 +127,13 @@ def test_parsearg():
             tctx.master.commands, "foo", command.Cuts
         ) == [["test"]]
 
+        assert command.parsearg(
+            tctx.master.commands, "foo", typing.Sequence[str]
+        ) == ["foo"]
+        assert command.parsearg(
+            tctx.master.commands, "foo, bar", typing.Sequence[str]
+        ) == ["foo", "bar"]
+
 
 class TDec:
     @command.command("cmd1")
diff --git a/test/mitmproxy/utils/test_typecheck.py b/test/mitmproxy/utils/test_typecheck.py
index 17f70d37..fe33070e 100644
--- a/test/mitmproxy/utils/test_typecheck.py
+++ b/test/mitmproxy/utils/test_typecheck.py
@@ -88,25 +88,26 @@ def test_check_any():
     typecheck.check_option_type("foo", None, typing.Any)
 
 
-def test_check_command_return_type():
-    assert(typecheck.check_command_return_type("foo", str))
-    assert(typecheck.check_command_return_type(["foo"], typing.Sequence[str]))
-    assert(typecheck.check_command_return_type(None, None))
-    assert(not typecheck.check_command_return_type(["foo"], typing.Sequence[int]))
-    assert(not typecheck.check_command_return_type("foo", typing.Sequence[int]))
-    assert(typecheck.check_command_return_type([["foo", b"bar"]], command.Cuts))
-    assert(not typecheck.check_command_return_type(["foo", b"bar"], command.Cuts))
-    assert(not typecheck.check_command_return_type([["foo", 22]], command.Cuts))
+def test_check_command_type():
+    assert(typecheck.check_command_type("foo", str))
+    assert(typecheck.check_command_type(["foo"], typing.Sequence[str]))
+    assert(not typecheck.check_command_type(["foo", 1], typing.Sequence[str]))
+    assert(typecheck.check_command_type(None, None))
+    assert(not typecheck.check_command_type(["foo"], typing.Sequence[int]))
+    assert(not typecheck.check_command_type("foo", typing.Sequence[int]))
+    assert(typecheck.check_command_type([["foo", b"bar"]], command.Cuts))
+    assert(not typecheck.check_command_type(["foo", b"bar"], command.Cuts))
+    assert(not typecheck.check_command_type([["foo", 22]], command.Cuts))
 
     # Python 3.5 only defines __parameters__
     m = mock.Mock()
     m.__str__ = lambda self: "typing.Sequence"
     m.__parameters__ = (int,)
 
-    typecheck.check_command_return_type([10], m)
+    typecheck.check_command_type([10], m)
 
     # Python 3.5 only defines __union_params__
     m = mock.Mock()
     m.__str__ = lambda self: "typing.Union"
     m.__union_params__ = (int,)
-    assert not typecheck.check_command_return_type([22], m)
+    assert not typecheck.check_command_type([22], m)