introduce @command.argument

This makes it possible to specify more specific type annotations at runtime,
so that both mypy and our command system are happy. The .argument(name, type=)
syntax is similar to click's, so it should be fairly extensible if we need it.

4 files changed, 40 insertions, 44 deletions

diff --git a/mitmproxy/addons/core.py b/mitmproxy/addons/core.py
index 69df006f..4191d490 100644
--- a/mitmproxy/addons/core.py
+++ b/mitmproxy/addons/core.py
@@ -8,13 +8,6 @@ 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:
@@ -103,10 +96,11 @@ class Core:
         ]
 
     @command.command("flow.set")
+    @command.argument("spec", type=command.Choice("flow.set.options"))
     def flow_set(
         self,
         flows: typing.Sequence[flow.Flow],
-        spec: FlowSetChoice,
+        spec: str,
         sval: str
     ) -> None:
         """
@@ -193,11 +187,12 @@ class Core:
         ctx.log.alert("Toggled encoding on %s flows." % len(updated))
 
     @command.command("flow.encode")
+    @command.argument("enc", type=command.Choice("flow.encode.options"))
     def encode(
         self,
         flows: typing.Sequence[flow.Flow],
         part: str,
-        enc: FlowEncodeChoice,
+        enc: str,
     ) -> None:
         """
             Encode flows with a specified encoding.
diff --git a/mitmproxy/command.py b/mitmproxy/command.py
index a14c95d0..b8de6e36 100644
--- a/mitmproxy/command.py
+++ b/mitmproxy/command.py
@@ -2,6 +2,7 @@
     This module manges and invokes typed commands.
 """
 import inspect
+import types
 import typing
 import shlex
 import textwrap
@@ -18,25 +19,6 @@ 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 = ""
-
-
 Path = typing.NewType("Path", str)
 
 
@@ -45,7 +27,7 @@ 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 hasattr(t, "options_command"):
+    if isinstance(t, Choice):
         return "choice"
     elif t == typing.Sequence[flow.Flow]:
         return "[flow]" if ret else "flowspec"
@@ -112,11 +94,11 @@ class Command:
             args = args[:len(self.paramtypes) - 1]
 
         pargs = []
-        for i in range(len(args)):
-            if typecheck.check_command_type(args[i], self.paramtypes[i]):
-                pargs.append(args[i])
+        for arg, paramtype in zip(args, self.paramtypes):
+            if typecheck.check_command_type(arg, paramtype):
+                pargs.append(arg)
             else:
-                pargs.append(parsearg(self.manager, args[i], self.paramtypes[i]))
+                pargs.append(parsearg(self.manager, arg, paramtype))
 
         if remainder:
             chk = typecheck.check_command_type(
@@ -183,8 +165,8 @@ 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")
+    if isinstance(argtype, Choice):
+        cmd = argtype.options_command
         opts = manager.call(cmd)
         if spec not in opts:
             raise exceptions.CommandError(
@@ -243,3 +225,26 @@ def command(path):
         wrapper.__dict__["command_path"] = path
         return wrapper
     return decorator
+
+
+class Choice:
+    def __init__(self, options_command):
+        self.options_command = options_command
+
+    def __instancecheck__(self, instance):
+        # return false here so that arguments are piped through parsearg,
+        # which does extended validation.
+        return False
+
+
+def argument(name, type):
+    """
+    Set the type of a command argument at runtime.
+    This is useful for more specific types such as command.Choice, which we cannot annotate
+    directly as mypy does not like that.
+    """
+    def decorator(f: types.FunctionType) -> types.FunctionType:
+        assert name in f.__annotations__
+        f.__annotations__[name] = type
+        return f
+    return decorator
diff --git a/mitmproxy/tools/console/consoleaddons.py b/mitmproxy/tools/console/consoleaddons.py
index 7772545d..471e3a53 100644
--- a/mitmproxy/tools/console/consoleaddons.py
+++ b/mitmproxy/tools/console/consoleaddons.py
@@ -31,12 +31,6 @@ 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):
@@ -363,7 +357,8 @@ class ConsoleAddon:
         ]
 
     @command.command("console.edit.focus")
-    def edit_focus(self, part: FocusChoice) -> None:
+    @command.argument("part", type=command.Choice("console.edit.focus.options"))
+    def edit_focus(self, part: str) -> None:
         """
             Edit a component of the currently focused flow.
         """
@@ -436,7 +431,8 @@ class ConsoleAddon:
         self._grideditor().cmd_spawn_editor()
 
     @command.command("console.flowview.mode.set")
-    def flowview_mode_set(self, mode: FlowViewModeChoice) -> None:
+    @command.argument("mode", type=command.Choice("console.flowview.mode.options"))
+    def flowview_mode_set(self, mode: str) -> None:
         """
             Set the display mode for the current flow view.
         """
diff --git a/mitmproxy/utils/typecheck.py b/mitmproxy/utils/typecheck.py
index c5e289a4..87a0e804 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(typeinfo, type)) or (not isinstance(value, typeinfo)):
+    elif not isinstance(value, typeinfo):
         return False
     return True