commands: refactor types

The type system was scattered over a number of places, making it hard to
follow. This collects all command types in types.py, and completion, validation
and parsing for each type is centralised. We should use the same mechanism for
options.

5 files changed, 200 insertions, 62 deletions

diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py
index 50ad3d55..f9315dd2 100644
--- a/test/mitmproxy/test_command.py
+++ b/test/mitmproxy/test_command.py
@@ -4,6 +4,7 @@ from mitmproxy import flow
 from mitmproxy import exceptions
 from mitmproxy.test import tflow
 from mitmproxy.test import taddons
+import mitmproxy.types
 import io
 import pytest
 
@@ -25,7 +26,7 @@ class TAddon:
         return foo
 
     @command.command("subcommand")
-    def subcommand(self, cmd: command.Cmd, *args: command.Arg) -> str:
+    def subcommand(self, cmd: mitmproxy.types.Cmd, *args: mitmproxy.types.Arg) -> str:
         return "ok"
 
     @command.command("empty")
@@ -39,12 +40,12 @@ class TAddon:
     def choices(self) -> typing.Sequence[str]:
         return ["one", "two", "three"]
 
-    @command.argument("arg", type=command.Choice("choices"))
+    @command.argument("arg", type=mitmproxy.types.Choice("choices"))
     def choose(self, arg: str) -> typing.Sequence[str]:
         return ["one", "two", "three"]
 
     @command.command("path")
-    def path(self, arg: command.Path) -> None:
+    def path(self, arg: mitmproxy.types.Path) -> None:
         pass
 
 
@@ -79,45 +80,45 @@ class TestCommand:
             [
                 "foo bar",
                 [
-                    command.ParseResult(value = "foo", type = command.Cmd),
+                    command.ParseResult(value = "foo", type = mitmproxy.types.Cmd),
                     command.ParseResult(value = "bar", type = str)
                 ],
             ],
             [
                 "foo 'bar",
                 [
-                    command.ParseResult(value = "foo", type = command.Cmd),
+                    command.ParseResult(value = "foo", type = mitmproxy.types.Cmd),
                     command.ParseResult(value = "'bar", type = str)
                 ]
             ],
-            ["a", [command.ParseResult(value = "a", type = command.Cmd)]],
-            ["", [command.ParseResult(value = "", type = command.Cmd)]],
+            ["a", [command.ParseResult(value = "a", type = mitmproxy.types.Cmd)]],
+            ["", [command.ParseResult(value = "", type = mitmproxy.types.Cmd)]],
             [
                 "cmd3 1",
                 [
-                    command.ParseResult(value = "cmd3", type = command.Cmd),
+                    command.ParseResult(value = "cmd3", type = mitmproxy.types.Cmd),
                     command.ParseResult(value = "1", type = int),
                 ]
             ],
             [
                 "cmd3 ",
                 [
-                    command.ParseResult(value = "cmd3", type = command.Cmd),
+                    command.ParseResult(value = "cmd3", type = mitmproxy.types.Cmd),
                     command.ParseResult(value = "", type = int),
                 ]
             ],
             [
                 "subcommand ",
                 [
-                    command.ParseResult(value = "subcommand", type = command.Cmd),
-                    command.ParseResult(value = "", type = command.Cmd),
+                    command.ParseResult(value = "subcommand", type = mitmproxy.types.Cmd),
+                    command.ParseResult(value = "", type = mitmproxy.types.Cmd),
                 ]
             ],
             [
                 "subcommand cmd3 ",
                 [
-                    command.ParseResult(value = "subcommand", type = command.Cmd),
-                    command.ParseResult(value = "cmd3", type = command.Cmd),
+                    command.ParseResult(value = "subcommand", type = mitmproxy.types.Cmd),
+                    command.ParseResult(value = "cmd3", type = mitmproxy.types.Cmd),
                     command.ParseResult(value = "", type = int),
                 ]
             ],
@@ -154,15 +155,15 @@ def test_typename():
     assert command.typename(str) == "str"
     assert command.typename(typing.Sequence[flow.Flow]) == "[flow]"
 
-    assert command.typename(command.Cuts) == "[cuts]"
-    assert command.typename(typing.Sequence[command.Cut]) == "[cut]"
+    assert command.typename(mitmproxy.types.Data) == "[data]"
+    assert command.typename(mitmproxy.types.CutSpec) == "[cut]"
 
     assert command.typename(flow.Flow) == "flow"
     assert command.typename(typing.Sequence[str]) == "[str]"
 
-    assert command.typename(command.Choice("foo")) == "choice"
-    assert command.typename(command.Path) == "path"
-    assert command.typename(command.Cmd) == "cmd"
+    assert command.typename(mitmproxy.types.Choice("foo")) == "choice"
+    assert command.typename(mitmproxy.types.Path) == "path"
+    assert command.typename(mitmproxy.types.Cmd) == "cmd"
 
 
 class DummyConsole:
@@ -172,7 +173,7 @@ class DummyConsole:
         return [tflow.tflow(resp=True)] * n
 
     @command.command("cut")
-    def cut(self, spec: str) -> command.Cuts:
+    def cut(self, spec: str) -> mitmproxy.types.Data:
         return [["test"]]
 
 
@@ -202,10 +203,6 @@ def test_parsearg():
             command.parsearg(tctx.master.commands, "foo", Exception)
 
         assert command.parsearg(
-            tctx.master.commands, "foo", command.Cuts
-        ) == [["test"]]
-
-        assert command.parsearg(
             tctx.master.commands, "foo", typing.Sequence[str]
         ) == ["foo"]
         assert command.parsearg(
@@ -215,18 +212,18 @@ def test_parsearg():
         a = TAddon()
         tctx.master.commands.add("choices", a.choices)
         assert command.parsearg(
-            tctx.master.commands, "one", command.Choice("choices"),
+            tctx.master.commands, "one", mitmproxy.types.Choice("choices"),
         ) == "one"
         with pytest.raises(exceptions.CommandError):
             assert command.parsearg(
-                tctx.master.commands, "invalid", command.Choice("choices"),
+                tctx.master.commands, "invalid", mitmproxy.types.Choice("choices"),
             )
 
         assert command.parsearg(
-            tctx.master.commands, "foo", command.Path
+            tctx.master.commands, "foo", mitmproxy.types.Path
         ) == "foo"
         assert command.parsearg(
-            tctx.master.commands, "foo", command.Cmd
+            tctx.master.commands, "foo", mitmproxy.types.Cmd
         ) == "foo"
 
 
@@ -272,5 +269,5 @@ def test_choice():
     basic typechecking for choices should fail as we cannot verify if strings are a valid choice
     at this point.
     """
-    c = command.Choice("foo")
+    c = mitmproxy.types.Choice("foo")
     assert not typecheck.check_command_type("foo", c)
diff --git a/test/mitmproxy/test_typemanager.py b/test/mitmproxy/test_typemanager.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/test/mitmproxy/test_typemanager.py
diff --git a/test/mitmproxy/test_types.py b/test/mitmproxy/test_types.py
new file mode 100644
index 00000000..81aaed74
--- /dev/null
+++ b/test/mitmproxy/test_types.py
@@ -0,0 +1,175 @@
+import pytest
+import os
+import typing
+import contextlib
+
+from mitmproxy.test import tutils
+import mitmproxy.exceptions
+import mitmproxy.types
+from mitmproxy.test import taddons
+from mitmproxy.test import tflow
+from mitmproxy import command
+from mitmproxy import flow
+
+from . import test_command
+
+
+@contextlib.contextmanager
+def chdir(path: str):
+    old_dir = os.getcwd()
+    os.chdir(path)
+    yield
+    os.chdir(old_dir)
+
+
+def test_bool():
+    with taddons.context() as tctx:
+        b = mitmproxy.types.Bool()
+        assert b.completion(tctx.master.commands, bool, "b") == ["false", "true"]
+        assert b.parse(tctx.master.commands, bool, "true") is True
+        assert b.parse(tctx.master.commands, bool, "false") is False
+        with pytest.raises(mitmproxy.exceptions.TypeError):
+            b.parse(tctx.master.commands, bool, "foo")
+
+
+def test_str():
+    with taddons.context() as tctx:
+        b = mitmproxy.types.Str()
+        assert b.completion(tctx.master.commands, str, "") == []
+        assert b.parse(tctx.master.commands, str, "foo") == "foo"
+
+
+def test_int():
+    with taddons.context() as tctx:
+        b = mitmproxy.types.Int()
+        assert b.completion(tctx.master.commands, int, "b") == []
+        assert b.parse(tctx.master.commands, int, "1") == 1
+        assert b.parse(tctx.master.commands, int, "999") == 999
+        with pytest.raises(mitmproxy.exceptions.TypeError):
+            b.parse(tctx.master.commands, int, "foo")
+
+
+def test_path():
+    with taddons.context() as tctx:
+        b = mitmproxy.types.PathType()
+        assert b.parse(tctx.master.commands, mitmproxy.types.Path, "/foo") == "/foo"
+        assert b.parse(tctx.master.commands, mitmproxy.types.Path, "/bar") == "/bar"
+
+        def normPathOpts(prefix, match):
+            ret = []
+            for s in b.completion(tctx.master.commands, mitmproxy.types.Path, match):
+                s = s[len(prefix):]
+                s = s.replace(os.sep, "/")
+                ret.append(s)
+            return ret
+
+        cd = os.path.normpath(tutils.test_data.path("mitmproxy/completion"))
+        assert normPathOpts(cd, cd) == ['/aaa', '/aab', '/aac', '/bbb/']
+        assert normPathOpts(cd, os.path.join(cd, "a")) == ['/aaa', '/aab', '/aac']
+        with chdir(cd):
+            assert normPathOpts("", "./") == ['./aaa', './aab', './aac', './bbb/']
+            assert normPathOpts("", "") == ['./aaa', './aab', './aac', './bbb/']
+        assert b.completion(
+            tctx.master.commands, mitmproxy.types.Path, "nonexistent"
+        ) == ["nonexistent"]
+
+
+def test_cmd():
+    with taddons.context() as tctx:
+        tctx.master.addons.add(test_command.TAddon())
+        b = mitmproxy.types.CmdType()
+        assert b.parse(tctx.master.commands, mitmproxy.types.Cmd, "foo") == "foo"
+        assert len(
+            b.completion(tctx.master.commands, mitmproxy.types.Cmd, "")
+        ) == len(tctx.master.commands.commands.keys())
+
+
+def test_cutspec():
+    with taddons.context() as tctx:
+        b = mitmproxy.types.CutSpecType()
+        b.parse(tctx.master.commands, mitmproxy.types.CutSpec, "foo,bar") == ["foo", "bar"]
+        assert b.completion(
+            tctx.master.commands, mitmproxy.types.CutSpec, "request.p"
+        ) == b.valid_prefixes
+        ret = b.completion(tctx.master.commands, mitmproxy.types.CutSpec, "request.port,f")
+        assert ret[0].startswith("request.port,")
+        assert len(ret) == len(b.valid_prefixes)
+
+
+def test_arg():
+    with taddons.context() as tctx:
+        b = mitmproxy.types.ArgType()
+        assert b.completion(tctx.master.commands, mitmproxy.types.Arg, "") == []
+        assert b.parse(tctx.master.commands, mitmproxy.types.Arg, "foo") == "foo"
+
+
+def test_strseq():
+    with taddons.context() as tctx:
+        b = mitmproxy.types.StrSeq()
+        assert b.completion(tctx.master.commands, typing.Sequence[str], "") == []
+        assert b.parse(tctx.master.commands, typing.Sequence[str], "foo") == ["foo"]
+        assert b.parse(tctx.master.commands, typing.Sequence[str], "foo,bar") == ["foo", "bar"]
+
+
+class DummyConsole:
+    @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) -> mitmproxy.types.Data:
+        return [["test"]]
+
+    @command.command("options")
+    def options(self) -> typing.Sequence[str]:
+        return ["one", "two", "three"]
+
+
+def test_flow():
+    with taddons.context() as tctx:
+        tctx.master.addons.add(DummyConsole())
+        b = mitmproxy.types.FlowType()
+        assert len(b.completion(tctx.master.commands, flow.Flow, "")) == len(b.valid_prefixes)
+        assert b.parse(tctx.master.commands, flow.Flow, "1")
+        with pytest.raises(mitmproxy.exceptions.TypeError):
+            assert b.parse(tctx.master.commands, flow.Flow, "0")
+        with pytest.raises(mitmproxy.exceptions.TypeError):
+            assert b.parse(tctx.master.commands, flow.Flow, "2")
+
+
+def test_flows():
+    with taddons.context() as tctx:
+        tctx.master.addons.add(DummyConsole())
+        b = mitmproxy.types.FlowsType()
+        assert len(
+            b.completion(tctx.master.commands, typing.Sequence[flow.Flow], "")
+        ) == len(b.valid_prefixes)
+        assert len(b.parse(tctx.master.commands, typing.Sequence[flow.Flow], "0")) == 0
+        assert len(b.parse(tctx.master.commands, typing.Sequence[flow.Flow], "1")) == 1
+        assert len(b.parse(tctx.master.commands, typing.Sequence[flow.Flow], "2")) == 2
+
+
+def test_data():
+    with taddons.context() as tctx:
+        b = mitmproxy.types.DataType()
+        with pytest.raises(mitmproxy.exceptions.TypeError):
+            b.parse(tctx.master.commands, mitmproxy.types.Data, "foo")
+        with pytest.raises(mitmproxy.exceptions.TypeError):
+            b.parse(tctx.master.commands, mitmproxy.types.Data, "foo")
+
+
+def test_choice():
+    with taddons.context() as tctx:
+        tctx.master.addons.add(DummyConsole())
+        b = mitmproxy.types.ChoiceType()
+        comp = b.completion(tctx.master.commands, mitmproxy.types.Choice("options"), "")
+        assert comp == ["one", "two", "three"]
+        assert b.parse(tctx.master.commands, mitmproxy.types.Choice("options"), "one") == "one"
+        with pytest.raises(mitmproxy.exceptions.TypeError):
+            b.parse(tctx.master.commands, mitmproxy.types.Choice("options"), "invalid")
+
+
+def test_typemanager():
+    assert mitmproxy.types.CommandTypes.get(bool, None)
+    assert mitmproxy.types.CommandTypes.get(mitmproxy.types.Choice("choide"), None)
diff --git a/test/mitmproxy/tools/console/test_commander.py b/test/mitmproxy/tools/console/test_commander.py
index 823af06d..34062dcb 100644
--- a/test/mitmproxy/tools/console/test_commander.py
+++ b/test/mitmproxy/tools/console/test_commander.py
@@ -1,36 +1,6 @@
-import os
-import contextlib
 
 from mitmproxy.tools.console.commander import commander
 from mitmproxy.test import taddons
-from mitmproxy.test import tutils
-
-
-@contextlib.contextmanager
-def chdir(path: str):
-    old_dir = os.getcwd()
-    os.chdir(path)
-    yield
-    os.chdir(old_dir)
-
-
-def normPathOpts(prefix, match):
-    ret = []
-    for s in commander.pathOptions(match):
-        s = s[len(prefix):]
-        s = s.replace(os.sep, "/")
-        ret.append(s)
-    return ret
-
-
-def test_pathOptions():
-    cd = os.path.normpath(tutils.test_data.path("mitmproxy/completion"))
-    assert normPathOpts(cd, cd) == ['/aaa', '/aab', '/aac', '/bbb/']
-    assert normPathOpts(cd, os.path.join(cd, "a")) == ['/aaa', '/aab', '/aac']
-    with chdir(cd):
-        assert normPathOpts("", "./") == ['./aaa', './aab', './aac', './bbb/']
-        assert normPathOpts("", "") == ['./aaa', './aab', './aac', './bbb/']
-    assert commander.pathOptions("nonexistent") == ["nonexistent"]
 
 
 class TestListCompleter:
diff --git a/test/mitmproxy/utils/test_typecheck.py b/test/mitmproxy/utils/test_typecheck.py
index 66b1884e..365509f1 100644
--- a/test/mitmproxy/utils/test_typecheck.py
+++ b/test/mitmproxy/utils/test_typecheck.py
@@ -4,7 +4,6 @@ from unittest import mock
 import pytest
 
 from mitmproxy.utils import typecheck
-from mitmproxy import command
 
 
 class TBase:
@@ -95,9 +94,6 @@ def test_check_command_type():
     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()