more robust argparse and arglist decorators

2 files changed, 62 insertions, 18 deletions

diff --git a/cmd2.py b/cmd2.py
index 0fc6c32c..e405bef9 100755
--- a/cmd2.py
+++ b/cmd2.py
@@ -242,6 +242,22 @@ def strip_quotes(arg):
     return arg
 
 
+def parse_quoted_string(cmdline):
+    """Parse a quoted string into a list of arguments."""
+    if isinstance(cmdline, list):
+        # arguments are already a list, return the list we were passed
+        lexed_arglist = cmdline
+    else:
+        # Use shlex to split the command line into a list of arguments based on shell rules
+        lexed_arglist = shlex.split(cmdline, posix=POSIX_SHLEX)
+        # If not using POSIX shlex, make sure to strip off outer quotes for convenience
+        if not POSIX_SHLEX and STRIP_QUOTES_FOR_NON_POSIX:
+            temp_arglist = []
+            for arg in lexed_arglist:
+                temp_arglist.append(strip_quotes(arg))
+            lexed_arglist = temp_arglist
+    return lexed_arglist
+
 def with_argument_parser(argparser):
     """A decorator to alter a cmd2 method to populate its ``opts``
     argument by parsing arguments with the given instance of
@@ -249,14 +265,7 @@ def with_argument_parser(argparser):
     """
     def arg_decorator(func):
         def cmd_wrapper(instance, cmdline):
-            # Use shlex to split the command line into a list of arguments based on shell rules
-            lexed_arglist = shlex.split(cmdline, posix=POSIX_SHLEX)
-            # If not using POSIX shlex, make sure to strip off outer quotes for convenience
-            if not POSIX_SHLEX and STRIP_QUOTES_FOR_NON_POSIX:
-                temp_arglist = []
-                for arg in lexed_arglist:
-                    temp_arglist.append(strip_quotes(arg))
-                lexed_arglist = temp_arglist
+            lexed_arglist = parse_quoted_string(cmdline)
             opts = argparser.parse_args(lexed_arglist)
             func(instance, lexed_arglist, opts)
 
@@ -282,14 +291,7 @@ def with_argument_list(func):
     With this decorator, the decorated method will receive a list
     of arguments parsed from user input using shlex.split()."""
     def cmd_wrapper(self, cmdline):
-        # Use shlex to split the command line into a list of arguments based on shell rules
-        lexed_arglist = shlex.split(cmdline, posix=POSIX_SHLEX)
-        # If not using POSIX shlex, make sure to strip off outer quotes for convenience
-        if not POSIX_SHLEX and STRIP_QUOTES_FOR_NON_POSIX:
-            temp_arglist = []
-            for arg in lexed_arglist:
-                temp_arglist.append(strip_quotes(arg))
-            lexed_arglist = temp_arglist
+        lexed_arglist = parse_quoted_string(cmdline)
         func(self, lexed_arglist)
     return cmd_wrapper
 
diff --git a/tests/test_argparse.py b/tests/test_argparse.py
index 7d6d99de..79075fe9 100644
--- a/tests/test_argparse.py
+++ b/tests/test_argparse.py
@@ -64,6 +64,17 @@ class ArgparseApp(cmd2.Cmd):
         else:
             self.stdout.write('False')
 
+
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument('args', nargs='*')
+    @cmd2.with_argument_list
+    @cmd2.with_argument_parser(argparser)
+    def do_arglistandargparser(self, arglist, args=None):
+        if isinstance(arglist, list):
+            self.stdout.write(' '.join(arglist))
+        else:
+            self.stdout.write('False')
+
     @cmd2.with_argument_list
     def do_arglist(self, arglist):
         if isinstance(arglist, list):
@@ -71,17 +82,32 @@ class ArgparseApp(cmd2.Cmd):
         else:
             self.stdout.write('False')
 
+    @cmd2.with_argument_list
+    @cmd2.with_argument_list
+    def do_arglisttwice(self, arglist):
+        if isinstance(arglist, list):
+            self.stdout.write(' '.join(arglist))
+        else:
+            self.stdout.write('False')
+
+
 class ArglistApp(cmd2.Cmd):
     def __init__(self):
         self.use_argument_list = True
         cmd2.Cmd.__init__(self)
 
     def do_arglist(self, arglist):
+        """Print true if the arglist parameter is passed as a list."""
         if isinstance(arglist, list):
             self.stdout.write('True')
         else:
             self.stdout.write('False')
 
+    @cmd2.with_argument_list
+    def do_arglistwithdecorator(self, arglist):
+        self.stdout.write(' '.join(arglist))
+        
+
 @pytest.fixture
 def argparse_app():
     app = ArgparseApp()
@@ -142,9 +168,25 @@ def test_argparse_arglist(argparse_app):
     assert out[0] == 'True'
 
 def test_arglist(argparse_app):
-    out = run_cmd(argparse_app, 'arglist "we should" get these in a list, not a string')
+    out = run_cmd(argparse_app, 'arglist "we  should" get these')
     assert out[0] == 'True'
 
+def test_arglist_decorator_twice(argparse_app):
+    out = run_cmd(argparse_app, 'arglisttwice "we  should" get these')
+    assert out[0] == 'we  should get these'
+
+def test_arglist_and_argparser(argparse_app):
+    out = run_cmd(argparse_app, 'arglistandargparser some "quoted   words"')
+    assert out[0] == 'some quoted   words'
+
 def test_use_argument_list(arglist_app):
-    out = run_cmd(arglist_app, 'arglist "we should" get these in a list, not a string')
+    out = run_cmd(arglist_app, 'arglist "we  should" get these in a list, not a string')
     assert out[0] == 'True'
+
+def test_arglist_attribute_and_decorator(arglist_app):
+    out = run_cmd(arglist_app, 'arglistwithdecorator "we  should" get these')
+    assert out[0] == 'we  should get these'
+
+#def test_arglist_help(arglist_app):
+#    out = run_cmd(arglist_app, 'help arglist')
+#    assert out[0] == 'True'