Merge pull request #2671 from mhils/command-argtypes

Introduce @command.argument

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.

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,26 +19,8 @@ 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)
+class Path(str):
+    pass
 
 
 def typename(t: type, ret: bool) -> str:
@@ -45,7 +28,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"
@@ -55,10 +38,8 @@ def typename(t: type, ret: bool) -> str:
         return "[cuts]" if ret else "cutspec"
     elif t == flow.Flow:
         return "flow"
-    elif t == Path:
-        return "path"
     elif issubclass(t, (str, int, bool)):
-        return t.__name__
+        return t.__name__.lower()
     else:  # pragma: no cover
         raise NotImplementedError(t)
 
@@ -112,11 +93,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,16 +164,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")
+    if isinstance(argtype, Choice):
+        cmd = 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 argtype == Path:
-        return spec
     elif issubclass(argtype, str):
         return spec
     elif argtype == bool:
@@ -243,3 +222,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

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.
         """

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
 

test/mitmproxy/test_command.py

@@ -7,9 +7,7 @@ from mitmproxy.test import taddons
 import io
 import pytest
 
-
-TChoice = typing.NewType("TChoice", command.Choice)
-TChoice.options_command = "choices"
+from mitmproxy.utils import typecheck
 
 
 class TAddon:
@@ -32,7 +30,8 @@ class TAddon:
     def choices(self) -> typing.Sequence[str]:
         return ["one", "two", "three"]
 
-    def choose(self, arg: TChoice) -> typing.Sequence[str]:  # type: ignore
+    @command.argument("arg", type=command.Choice("choices"))
+    def choose(self, arg: str) -> typing.Sequence[str]:
         return ["one", "two", "three"]
 
     def path(self, arg: command.Path) -> None:
@@ -99,7 +98,7 @@ def test_typename():
     assert command.typename(flow.Flow, False) == "flow"
     assert command.typename(typing.Sequence[str], False) == "[str]"
 
-    assert command.typename(TChoice, False) == "choice"
+    assert command.typename(command.Choice("foo"), False) == "choice"
     assert command.typename(command.Path, False) == "path"
 
 
@@ -153,11 +152,11 @@ def test_parsearg():
         a = TAddon()
         tctx.master.commands.add("choices", a.choices)
         assert command.parsearg(
-            tctx.master.commands, "one", TChoice,
+            tctx.master.commands, "one", command.Choice("choices"),
         ) == "one"
         with pytest.raises(exceptions.CommandError):
             assert command.parsearg(
-                tctx.master.commands, "invalid", TChoice,
+                tctx.master.commands, "invalid", command.Choice("choices"),
             )
 
         assert command.parsearg(
@@ -199,4 +198,13 @@ def test_verify_arg_signature():
     with pytest.raises(exceptions.CommandError):
         command.verify_arg_signature(lambda: None, [1, 2], {})
         print('hello there')
-    command.verify_arg_signature(lambda a, b: None, [1, 2], {})
+    command.verify_arg_signature(lambda a, b: None, [1, 2], {})
+
+
+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")
+    assert not typecheck.check_command_type("foo", c)