diff options
| -rw-r--r-- | mitmproxy/addons/core.py | 13 | ||||
| -rw-r--r-- | mitmproxy/command.py | 66 | ||||
| -rw-r--r-- | mitmproxy/tools/console/consoleaddons.py | 12 | ||||
| -rw-r--r-- | mitmproxy/utils/typecheck.py | 2 | ||||
| -rw-r--r-- | test/mitmproxy/test_command.py | 24 |
5 files changed, 59 insertions, 58 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..c4821973 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,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 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 diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py index abb09cea..e1879ba2 100644 --- a/test/mitmproxy/test_command.py +++ b/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], {})
\ No newline at end of file + 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) |