commands: formalise a Choice type

This resolves as a string during MyPy checks, but at runtime has an additional
attribute that is a command that returns valid options.

This is very ugly and clumsy, basically because MyPy is super restrictive about
what it accepts as a type. Almost any attempt to construct these types in a
more sophisticated way fails in one way or another. I'm open to suggestions.

7 files changed, 79 insertions, 25 deletions

diff --git a/mitmproxy/addons/core.py b/mitmproxy/addons/core.py
index 33d67279..8a63422d 100644
--- a/mitmproxy/addons/core.py
+++ b/mitmproxy/addons/core.py
@@ -8,6 +8,13 @@ from mitmproxy import optmanager
 from mitmproxy.net.http import status_codes
 
 
+FlowSetChoice = typing.NewType("FlowSetChoice", command.Choice)
+FlowSetChoice.options_command = "flow.set.options"
+
+FlowEncodeChoice = typing.NewType("FlowEncodeChoice", command.Choice)
+FlowEncodeChoice.options_command = "flow.encode.options"
+
+
 class Core:
     @command.command("set")
     def set(self, *spec: str) -> None:
@@ -98,17 +105,13 @@ class Core:
     @command.command("flow.set")
     def flow_set(
         self,
-        flows: typing.Sequence[flow.Flow], spec: str, sval: str
+        flows: typing.Sequence[flow.Flow],
+        spec: FlowSetChoice,
+        sval: str
     ) -> None:
         """
             Quickly set a number of common values on flows.
         """
-        opts = self.flow_set_options()
-        if spec not in opts:
-            raise exceptions.CommandError(
-                "Set spec must be one of: %s." % ", ".join(opts)
-            )
-
         val = sval  # type: typing.Union[int, str]
         if spec == "status_code":
             try:
@@ -190,13 +193,15 @@ class Core:
         ctx.log.alert("Toggled encoding on %s flows." % len(updated))
 
     @command.command("flow.encode")
-    def encode(self, flows: typing.Sequence[flow.Flow], part: str, enc: str) -> None:
+    def encode(
+        self,
+        flows: typing.Sequence[flow.Flow],
+        part: str,
+        enc: FlowEncodeChoice,
+    ) -> None:
         """
             Encode flows with a specified encoding.
         """
-        if enc not in self.encode_options():
-            raise exceptions.CommandError("Invalid encoding format: %s" % enc)
-
         updated = []
         for f in flows:
             p = getattr(f, part, None)
@@ -212,7 +217,6 @@ class Core:
     def encode_options(self) -> typing.Sequence[str]:
         """
             The possible values for an encoding specification.
-
         """
         return ["gzip", "deflate", "br"]
 
diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
index 8ae1f341..6f0fd131 100644
--- a/mitmproxy/addons/view.py
+++ b/mitmproxy/addons/view.py
@@ -238,7 +238,7 @@ class View(collections.Sequence):
     @command.command("view.order.options")
     def order_options(self) -> typing.Sequence[str]:
         """
-            A list of all the orders we support.
+            Choices supported by the console_order option.
         """
         return list(sorted(self.orders.keys()))
 
diff --git a/mitmproxy/command.py b/mitmproxy/command.py
index eae3d80c..25e00174 100644
--- a/mitmproxy/command.py
+++ b/mitmproxy/command.py
@@ -18,12 +18,33 @@ Cuts = typing.Sequence[
 ]
 
 
+# A str that is validated at runtime by calling a command that returns options.
+#
+# This requires some explanation. We want to construct a type with two aims: it
+# must be detected as str by mypy, and it has to be decorated at runtime with an
+# options_commmand attribute that tells us where to look up options for runtime
+# validation. Unfortunately, mypy is really, really obtuse about what it detects
+# as a type - any construction of these types at runtime barfs. The effect is
+# that while the annotation mechanism is very generaly, if you also use mypy
+# you're hamstrung. So the middle road is to declare a derived type, which is
+# then used very clumsily as follows:
+#
+#       MyType = typing.NewType("MyType", command.Choice)
+#       MyType.options_command = "my.command"
+#
+# The resulting type is then used in the function argument decorator.
+class Choice(str):
+    options_command = ""
+
+
 def typename(t: type, ret: bool) -> str:
     """
         Translates a type to an explanatory string. If ret is True, we're
         looking at a return type, else we're looking at a parameter type.
     """
-    if issubclass(t, (str, int, bool)):
+    if hasattr(t, "options_command"):
+        return "choice"
+    elif issubclass(t, (str, int, bool)):
         return t.__name__
     elif t == typing.Sequence[flow.Flow]:
         return "[flow]" if ret else "flowspec"
@@ -157,6 +178,14 @@ def parsearg(manager: CommandManager, spec: str, argtype: type) -> typing.Any:
     """
         Convert a string to a argument to the appropriate type.
     """
+    if hasattr(argtype, "options_command"):
+        cmd = getattr(argtype, "options_command")
+        opts = manager.call(cmd)
+        if spec not in opts:
+            raise exceptions.CommandError(
+                "Invalid choice: see %s for options" % cmd
+            )
+        return spec
     if issubclass(argtype, str):
         return spec
     elif argtype == bool:
diff --git a/mitmproxy/tools/console/consoleaddons.py b/mitmproxy/tools/console/consoleaddons.py
index 81923baa..4b0bb00d 100644
--- a/mitmproxy/tools/console/consoleaddons.py
+++ b/mitmproxy/tools/console/consoleaddons.py
@@ -31,6 +31,12 @@ console_layouts = [
     "horizontal",
 ]
 
+FocusChoice = typing.NewType("FocusChoice", command.Choice)
+FocusChoice.options_command = "console.edit.focus.options"
+
+FlowViewModeChoice = typing.NewType("FlowViewModeChoice", command.Choice)
+FlowViewModeChoice.options_command = "console.flowview.mode.options"
+
 
 class Logger:
     def log(self, evt):
@@ -111,8 +117,7 @@ class ConsoleAddon:
     @command.command("console.layout.options")
     def layout_options(self) -> typing.Sequence[str]:
         """
-            Returns the valid options for console layout. Use these by setting
-            the console_layout option.
+            Returns the available options for the consoler_layout option.
         """
         return ["single", "vertical", "horizontal"]
 
@@ -358,7 +363,7 @@ class ConsoleAddon:
         ]
 
     @command.command("console.edit.focus")
-    def edit_focus(self, part: str) -> None:
+    def edit_focus(self, part: FocusChoice) -> None:
         """
             Edit a component of the currently focused flow.
         """
@@ -431,7 +436,7 @@ class ConsoleAddon:
         self._grideditor().cmd_spawn_editor()
 
     @command.command("console.flowview.mode.set")
-    def flowview_mode_set(self, mode: str) -> None:
+    def flowview_mode_set(self, mode: FlowViewModeChoice) -> None:
         """
             Set the display mode for the current flow view.
         """
diff --git a/mitmproxy/utils/typecheck.py b/mitmproxy/utils/typecheck.py
index 87a0e804..c5e289a4 100644
--- a/mitmproxy/utils/typecheck.py
+++ b/mitmproxy/utils/typecheck.py
@@ -31,7 +31,7 @@ def check_command_type(value: typing.Any, typeinfo: typing.Any) -> bool:
                 return False
     elif value is None and typeinfo is None:
         return True
-    elif not isinstance(value, typeinfo):
+    elif (not isinstance(typeinfo, type)) or (not isinstance(value, typeinfo)):
         return False
     return True
 
diff --git a/test/mitmproxy/addons/test_core.py b/test/mitmproxy/addons/test_core.py
index c132d80a..5aa4ef37 100644
--- a/test/mitmproxy/addons/test_core.py
+++ b/test/mitmproxy/addons/test_core.py
@@ -69,9 +69,6 @@ def test_flow_set():
         f = tflow.tflow(resp=True)
         assert sa.flow_set_options()
 
-        with pytest.raises(exceptions.CommandError):
-            sa.flow_set([f], "flibble", "post")
-
         assert f.request.method != "post"
         sa.flow_set([f], "method", "post")
         assert f.request.method == "POST"
@@ -126,9 +123,6 @@ def test_encoding():
         sa.encode_toggle([f], "request")
         assert "content-encoding" not in f.request.headers
 
-        with pytest.raises(exceptions.CommandError):
-            sa.encode([f], "request", "invalid")
-
 
 def test_options(tmpdir):
     p = str(tmpdir.join("path"))
diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py
index 43b97742..cb9dc4ed 100644
--- a/test/mitmproxy/test_command.py
+++ b/test/mitmproxy/test_command.py
@@ -8,6 +8,10 @@ import io
 import pytest
 
 
+TChoice = typing.NewType("TChoice", command.Choice)
+TChoice.options_command = "choices"
+
+
 class TAddon:
     def cmd1(self, foo: str) -> str:
         """cmd1 help"""
@@ -25,6 +29,12 @@ class TAddon:
     def varargs(self, one: str, *var: str) -> typing.Sequence[str]:
         return list(var)
 
+    def choices(self) -> typing.Sequence[str]:
+        return ["one", "two", "three"]
+
+    def choose(self, arg: TChoice) -> typing.Sequence[str]:  # type: ignore
+        return ["one", "two", "three"]
+
 
 class TestCommand:
     def test_varargs(self):
@@ -86,6 +96,8 @@ def test_typename():
     assert command.typename(flow.Flow, False) == "flow"
     assert command.typename(typing.Sequence[str], False) == "[str]"
 
+    assert command.typename(TChoice, False) == "choice"
+
 
 class DummyConsole:
     @command.command("view.resolve")
@@ -134,6 +146,16 @@ def test_parsearg():
             tctx.master.commands, "foo, bar", typing.Sequence[str]
         ) == ["foo", "bar"]
 
+        a = TAddon()
+        tctx.master.commands.add("choices", a.choices)
+        assert command.parsearg(
+            tctx.master.commands, "one", TChoice,
+        ) == "one"
+        with pytest.raises(exceptions.CommandError):
+            assert command.parsearg(
+                tctx.master.commands, "invalid", TChoice,
+            )
+
 
 class TDec:
     @command.command("cmd1")