From 38b37ba7f51c72bb94e1a5d2c9c7cf9836cb5a06 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Tue, 19 Dec 2017 07:38:21 +1300 Subject: [PATCH] types: cleanups and minor refactorings --- mitmproxy/command.py | 2 +- mitmproxy/tools/console/consoleaddons.py | 12 ++- mitmproxy/tools/console/defaultkeys.py | 2 +- mitmproxy/types.py | 116 +++++++++++------------ test/mitmproxy/test_types.py | 24 ++--- 5 files changed, 80 insertions(+), 76 deletions(-) diff --git a/mitmproxy/command.py b/mitmproxy/command.py index a77658fd4..e93a0cfa4 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -114,7 +114,7 @@ ParseResult = typing.NamedTuple( ) -class CommandManager: +class CommandManager(mitmproxy.types._CommandBase): def __init__(self, master): self.master = master self.commands = {} diff --git a/mitmproxy/tools/console/consoleaddons.py b/mitmproxy/tools/console/consoleaddons.py index 37647e60e..6ddcf394c 100644 --- a/mitmproxy/tools/console/consoleaddons.py +++ b/mitmproxy/tools/console/consoleaddons.py @@ -104,7 +104,7 @@ class ConsoleAddon: @command.command("console.layout.options") def layout_options(self) -> typing.Sequence[str]: """ - Returns the available options for the consoler_layout option. + Returns the available options for the console_layout option. """ return ["single", "vertical", "horizontal"] @@ -243,7 +243,11 @@ class ConsoleAddon: @command.command("console.choose.cmd") def console_choose_cmd( - self, prompt: str, choicecmd: mitmproxy.types.Cmd, *cmd: mitmproxy.types.Arg + self, + prompt: str, + choicecmd: mitmproxy.types.Cmd, + subcmd: mitmproxy.types.Cmd, + *args: mitmproxy.types.Arg ) -> None: """ Prompt the user to choose from a list of strings returned by a @@ -254,10 +258,10 @@ class ConsoleAddon: def callback(opt): # We're now outside of the call context... - repl = " ".join(cmd) + repl = " ".join(args) repl = repl.replace("{choice}", opt) try: - self.master.commands.call(repl) + self.master.commands.call(subcmd + " " + repl) except exceptions.CommandError as e: signals.status_message.send(message=str(e)) diff --git a/mitmproxy/tools/console/defaultkeys.py b/mitmproxy/tools/console/defaultkeys.py index 0fdec10c8..f8a3df2df 100644 --- a/mitmproxy/tools/console/defaultkeys.py +++ b/mitmproxy/tools/console/defaultkeys.py @@ -41,7 +41,7 @@ def map(km): "e", """ console.choose.cmd Format export.formats - console.command export.file {choice} @focus '' + console.command export.file {choice} @focus """, ["flowlist", "flowview"], "Export this flow to file" diff --git a/mitmproxy/types.py b/mitmproxy/types.py index e1b6f95db..b6b414ba8 100644 --- a/mitmproxy/types.py +++ b/mitmproxy/types.py @@ -40,39 +40,39 @@ class Choice: # annotations can cause circular dependencies where there were none before. # Rather than putting types and the CommandManger in the same file, we introduce # a stub type with the signature we use. -class _CommandStub: - commands = {} # type: typing.Mapping[str, typing.Any] +class _CommandBase: + commands = {} # type: typing.MutableMapping[str, typing.Any] - def call_args(self, path: str, args: typing.Sequence[str]) -> typing.Any: # pragma: no cover - pass + def call_args(self, path: str, args: typing.Sequence[str]) -> typing.Any: + raise NotImplementedError - def call(self, args: typing.Sequence[str]) -> typing.Any: # pragma: no cover - pass + def call(self, cmd: str) -> typing.Any: + raise NotImplementedError -class BaseType: +class _BaseType: typ = object # type: typing.Type display = "" # type: str def completion( - self, manager: _CommandStub, t: type, s: str + self, manager: _CommandBase, t: typing.Any, s: str ) -> typing.Sequence[str]: # pragma: no cover pass def parse( - self, manager: _CommandStub, t: type, s: str + self, manager: _CommandBase, t: typing.Any, s: str ) -> typing.Any: # pragma: no cover pass -class Bool(BaseType): +class _BoolType(_BaseType): typ = bool display = "bool" - def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]: + def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]: return ["false", "true"] - def parse(self, manager: _CommandStub, t: type, s: str) -> bool: + def parse(self, manager: _CommandBase, t: type, s: str) -> bool: if s == "true": return True elif s == "false": @@ -83,36 +83,36 @@ class Bool(BaseType): ) -class Str(BaseType): +class _StrType(_BaseType): typ = str display = "str" - def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]: + def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]: return [] - def parse(self, manager: _CommandStub, t: type, s: str) -> str: + def parse(self, manager: _CommandBase, t: type, s: str) -> str: return s -class Int(BaseType): +class _IntType(_BaseType): typ = int display = "int" - def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]: + def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]: return [] - def parse(self, manager: _CommandStub, t: type, s: str) -> int: + def parse(self, manager: _CommandBase, t: type, s: str) -> int: try: return int(s) except ValueError as e: raise exceptions.TypeError from e -class PathType(BaseType): +class _PathType(_BaseType): typ = Path display = "path" - def completion(self, manager: _CommandStub, t: type, start: str) -> typing.Sequence[str]: + def completion(self, manager: _CommandBase, t: type, start: str) -> typing.Sequence[str]: if not start: start = "./" path = os.path.expanduser(start) @@ -134,44 +134,44 @@ class PathType(BaseType): ret.sort() return ret - def parse(self, manager: _CommandStub, t: type, s: str) -> str: + def parse(self, manager: _CommandBase, t: type, s: str) -> str: return s -class CmdType(BaseType): +class _CmdType(_BaseType): typ = Cmd display = "cmd" - def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]: + def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]: return list(manager.commands.keys()) - def parse(self, manager: _CommandStub, t: type, s: str) -> str: + def parse(self, manager: _CommandBase, t: type, s: str) -> str: return s -class ArgType(BaseType): +class _ArgType(_BaseType): typ = Arg display = "arg" - def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]: + def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]: return [] - def parse(self, manager: _CommandStub, t: type, s: str) -> str: + def parse(self, manager: _CommandBase, t: type, s: str) -> str: return s -class StrSeq(BaseType): +class _StrSeqType(_BaseType): typ = typing.Sequence[str] display = "[str]" - def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]: + def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]: return [] - def parse(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]: + def parse(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]: return [x.strip() for x in s.split(",")] -class CutSpecType(BaseType): +class _CutSpecType(_BaseType): typ = CutSpec display = "[cut]" valid_prefixes = [ @@ -212,7 +212,7 @@ class CutSpecType(BaseType): "server_conn.ssl_established", ] - def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]: + def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]: spec = s.split(",") opts = [] for pref in self.valid_prefixes: @@ -220,12 +220,12 @@ class CutSpecType(BaseType): opts.append(",".join(spec)) return opts - def parse(self, manager: _CommandStub, t: type, s: str) -> CutSpec: + def parse(self, manager: _CommandBase, t: type, s: str) -> CutSpec: parts = s.split(",") # type: typing.Any return parts -class BaseFlowType(BaseType): +class _BaseFlowType(_BaseType): valid_prefixes = [ "@all", "@focus", @@ -248,15 +248,15 @@ class BaseFlowType(BaseType): "~c", ] - def completion(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[str]: + def completion(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[str]: return self.valid_prefixes -class FlowType(BaseFlowType): +class _FlowType(_BaseFlowType): typ = flow.Flow display = "flow" - def parse(self, manager: _CommandStub, t: type, s: str) -> flow.Flow: + def parse(self, manager: _CommandBase, t: type, s: str) -> flow.Flow: flows = manager.call_args("view.resolve", [s]) if len(flows) != 1: raise exceptions.TypeError( @@ -265,37 +265,37 @@ class FlowType(BaseFlowType): return flows[0] -class FlowsType(BaseFlowType): +class _FlowsType(_BaseFlowType): typ = typing.Sequence[flow.Flow] display = "[flow]" - def parse(self, manager: _CommandStub, t: type, s: str) -> typing.Sequence[flow.Flow]: + def parse(self, manager: _CommandBase, t: type, s: str) -> typing.Sequence[flow.Flow]: return manager.call_args("view.resolve", [s]) -class DataType: +class _DataType(_BaseType): typ = Data display = "[data]" def completion( - self, manager: _CommandStub, t: type, s: str + self, manager: _CommandBase, t: type, s: str ) -> typing.Sequence[str]: # pragma: no cover raise exceptions.TypeError("data cannot be passed as argument") def parse( - self, manager: _CommandStub, t: type, s: str + self, manager: _CommandBase, t: type, s: str ) -> typing.Any: # pragma: no cover raise exceptions.TypeError("data cannot be passed as argument") -class ChoiceType: +class _ChoiceType(_BaseType): typ = Choice display = "choice" - def completion(self, manager: _CommandStub, t: Choice, s: str) -> typing.Sequence[str]: + def completion(self, manager: _CommandBase, t: Choice, s: str) -> typing.Sequence[str]: return manager.call(t.options_command) - def parse(self, manager: _CommandStub, t: Choice, s: str) -> str: + def parse(self, manager: _CommandBase, t: Choice, s: str) -> str: opts = manager.call(t.options_command) if s not in opts: raise exceptions.TypeError("Invalid choice.") @@ -308,23 +308,23 @@ class TypeManager: for t in types: self.typemap[t.typ] = t() - def get(self, t: type, default=None) -> BaseType: + def get(self, t: type, default=None) -> _BaseType: if type(t) in self.typemap: return self.typemap[type(t)] return self.typemap.get(t, default) CommandTypes = TypeManager( - ArgType, - Bool, - ChoiceType, - CmdType, - CutSpecType, - DataType, - FlowType, - FlowsType, - Int, - PathType, - Str, - StrSeq, + _ArgType, + _BoolType, + _ChoiceType, + _CmdType, + _CutSpecType, + _DataType, + _FlowType, + _FlowsType, + _IntType, + _PathType, + _StrType, + _StrSeqType, ) diff --git a/test/mitmproxy/test_types.py b/test/mitmproxy/test_types.py index 81aaed749..900322046 100644 --- a/test/mitmproxy/test_types.py +++ b/test/mitmproxy/test_types.py @@ -24,7 +24,7 @@ def chdir(path: str): def test_bool(): with taddons.context() as tctx: - b = mitmproxy.types.Bool() + b = mitmproxy.types._BoolType() 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 @@ -34,14 +34,14 @@ def test_bool(): def test_str(): with taddons.context() as tctx: - b = mitmproxy.types.Str() + b = mitmproxy.types._StrType() 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() + b = mitmproxy.types._IntType() 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 @@ -51,7 +51,7 @@ def test_int(): def test_path(): with taddons.context() as tctx: - b = mitmproxy.types.PathType() + 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" @@ -77,7 +77,7 @@ def test_path(): def test_cmd(): with taddons.context() as tctx: tctx.master.addons.add(test_command.TAddon()) - b = mitmproxy.types.CmdType() + 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, "") @@ -86,7 +86,7 @@ def test_cmd(): def test_cutspec(): with taddons.context() as tctx: - b = mitmproxy.types.CutSpecType() + 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" @@ -98,14 +98,14 @@ def test_cutspec(): def test_arg(): with taddons.context() as tctx: - b = mitmproxy.types.ArgType() + 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() + b = mitmproxy.types._StrSeqType() 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"] @@ -129,7 +129,7 @@ class DummyConsole: def test_flow(): with taddons.context() as tctx: tctx.master.addons.add(DummyConsole()) - b = mitmproxy.types.FlowType() + 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): @@ -141,7 +141,7 @@ def test_flow(): def test_flows(): with taddons.context() as tctx: tctx.master.addons.add(DummyConsole()) - b = mitmproxy.types.FlowsType() + b = mitmproxy.types._FlowsType() assert len( b.completion(tctx.master.commands, typing.Sequence[flow.Flow], "") ) == len(b.valid_prefixes) @@ -152,7 +152,7 @@ def test_flows(): def test_data(): with taddons.context() as tctx: - b = mitmproxy.types.DataType() + 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): @@ -162,7 +162,7 @@ def test_data(): def test_choice(): with taddons.context() as tctx: tctx.master.addons.add(DummyConsole()) - b = mitmproxy.types.ChoiceType() + 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"