commands: add a parser for partial commands

We only return Cmd and str types for the moment.

2 files changed, 35 insertions, 0 deletions

diff --git a/mitmproxy/command.py b/mitmproxy/command.py
index a638b7ee..f73eeb68 100644
--- a/mitmproxy/command.py
+++ b/mitmproxy/command.py
@@ -3,6 +3,7 @@
 """
 import inspect
 import types
+import io
 import typing
 import shlex
 import textwrap
@@ -137,6 +138,30 @@ class CommandManager:
     def add(self, path: str, func: typing.Callable):
         self.commands[path] = Command(self, path, func)
 
+    def parse_partial(self, cmdstr: str) -> typing.Sequence[typing.Tuple[str, type]]:
+        """
+            Parse a possibly partial command. Return a sequence of (part, type) tuples.
+        """
+        parts: typing.List[typing.Tuple[str, type]] = []
+        buf = io.StringIO(cmdstr)
+        # mypy mis-identifies shlex.shlex as abstract
+        lex = shlex.shlex(buf)  # type: ignore
+        while 1:
+            remainder = cmdstr[buf.tell():]
+            try:
+                t = lex.get_token()
+            except ValueError:
+                parts.append((remainder, str))
+                break
+            if not t:
+                break
+            typ: type = str
+            # First value is a special case: it has to be a command
+            if not parts:
+                typ = Cmd
+            parts.append((t, typ))
+        return parts
+
     def call_args(self, path, args):
         """
             Call a command using a list of string arguments. May raise CommandError.
diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py
index c8007463..5218042c 100644
--- a/test/mitmproxy/test_command.py
+++ b/test/mitmproxy/test_command.py
@@ -64,6 +64,16 @@ class TestCommand:
             c = command.Command(cm, "cmd.three", a.cmd3)
             assert c.call(["1"]) == 1
 
+    def test_parse_partial(self):
+        tests = [
+            ["foo bar", [("foo", command.Cmd), ("bar", str)]],
+            ["foo 'bar", [("foo", command.Cmd), ("'bar", str)]],
+        ]
+        with taddons.context() as tctx:
+            cm = command.CommandManager(tctx.master)
+            for s, expected in tests:
+                assert cm.parse_partial(s) == expected
+
 
 def test_simple():
     with taddons.context() as tctx: