Merge pull request #636 from python-cmd2/with_argument_list

Fix bug in with_argument_list decorator

5 files changed, 33 insertions, 15 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index dd9966cc..fcf90db8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,7 @@
 ## 0.9.11 (TBD, 2019)
 * Bug Fixes
     * Fixed bug in how **history** command deals with multiline commands when output to a script
+    * Fixed a bug when the ``with_argument_list`` decorator is called with the optional ``preserve_quotes`` argument
 * Enhancements
     * Improvements to the **history** command
         * Simplified the display format and made it more similar to **bash**
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index 5912bee3..66de8473 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -176,25 +176,30 @@ def with_category(category: str) -> Callable:
     return cat_decorator
 
 
-def with_argument_list(func: Callable[[Statement], Optional[bool]],
-                       preserve_quotes: bool = False) -> Callable[[List], Optional[bool]]:
+def with_argument_list(*args: List[Callable], preserve_quotes: bool = False) -> Callable[[List], Optional[bool]]:
     """A decorator to alter the arguments passed to a do_* cmd2 method. Default passes a string of whatever the user
     typed. With this decorator, the decorated method will receive a list of arguments parsed from user input using
     shlex.split().
 
-    :param func: do_* method this decorator is wrapping
+    :param args: Single-element positional argument list containing do_* method this decorator is wrapping
     :param preserve_quotes: if True, then argument quotes will not be stripped
     :return: function that gets passed a list of argument strings
     """
     import functools
 
-    @functools.wraps(func)
-    def cmd_wrapper(self, cmdline):
-        lexed_arglist = parse_quoted_string(cmdline, preserve_quotes)
-        return func(self, lexed_arglist)
+    def arg_decorator(func: Callable):
+        @functools.wraps(func)
+        def cmd_wrapper(self, cmdline):
+            lexed_arglist = parse_quoted_string(cmdline, preserve_quotes)
+            return func(self, lexed_arglist)
 
-    cmd_wrapper.__doc__ = func.__doc__
-    return cmd_wrapper
+        cmd_wrapper.__doc__ = func.__doc__
+        return cmd_wrapper
+
+    if len(args) == 1 and callable(args[0]):
+        return arg_decorator(args[0])
+    else:
+        return arg_decorator
 
 
 def with_argparser_and_unknown_args(argparser: argparse.ArgumentParser, preserve_quotes: bool = False) -> \
diff --git a/docs/freefeatures.rst b/docs/freefeatures.rst
index 9e8c2cbf..5c246798 100644
--- a/docs/freefeatures.rst
+++ b/docs/freefeatures.rst
@@ -212,15 +212,14 @@ of using ``pyscript`` is shown below  along with the arg_printer_ script::
 
 .. note::
 
-    If you want to be able to pass arguments with spaces to scripts, then we strongly recommend using one of the decorators,
+    If you want to be able to pass arguments with spaces to commands, then we strongly recommend using one of the decorators,
     such as ``with_argument_list``.  ``cmd2`` will pass your **do_*** methods a list of arguments in this case.
 
-    When using this decorator, you can then put arguments in quotes like so (NOTE: the ``do_pyscript`` method uses this decorator::
+    When using this decorator, you can then put arguments in quotes like so::
 
-        (Cmd) pyscript examples/scripts/arg_printer.py hello '23 fnord'
-        Running Python script 'arg_printer.py' which was called with 2 arguments
-        arg 1: 'hello'
-        arg 2: '23 fnord'
+        $ examples/arg_print.py
+        (Cmd) lprint foo "bar baz"
+        lprint was called with the following list of arguments: ['foo', 'bar baz']
 
 .. _arg_printer: https://github.com/python-cmd2/cmd2/blob/master/examples/scripts/arg_printer.py
 
diff --git a/examples/arg_print.py b/examples/arg_print.py
index 1a103858..18d21787 100755
--- a/examples/arg_print.py
+++ b/examples/arg_print.py
@@ -38,6 +38,11 @@ class ArgumentAndOptionPrinter(cmd2.Cmd):
         """Print the argument list this basic command is called with."""
         self.poutput('lprint was called with the following list of arguments: {!r}'.format(arglist))
 
+    @cmd2.with_argument_list(preserve_quotes=True)
+    def do_rprint(self, arglist):
+        """Print the argument list this basic command is called with (with quotes preserved)."""
+        self.poutput('rprint was called with the following list of arguments: {!r}'.format(arglist))
+
     oprint_parser = argparse.ArgumentParser()
     oprint_parser.add_argument('-p', '--piglatin', action='store_true', help='atinLay')
     oprint_parser.add_argument('-s', '--shout', action='store_true', help='N00B EMULATION MODE')
diff --git a/tests/test_argparse.py b/tests/test_argparse.py
index 6b810b44..f5948f03 100644
--- a/tests/test_argparse.py
+++ b/tests/test_argparse.py
@@ -68,6 +68,10 @@ class ArgparseApp(cmd2.Cmd):
         else:
             self.stdout.write('False')
 
+    @cmd2.with_argument_list(preserve_quotes=True)
+    def do_preservelist(self, arglist):
+        self.stdout.write('{}'.format(arglist))
+
     @cmd2.with_argument_list
     @cmd2.with_argument_list
     def do_arglisttwice(self, arglist):
@@ -170,6 +174,10 @@ def test_arglist(argparse_app):
     out = run_cmd(argparse_app, 'arglist "we  should" get these')
     assert out[0] == 'True'
 
+def test_preservelist(argparse_app):
+    out = run_cmd(argparse_app, 'preservelist foo "bar baz"')
+    assert out[0] == "['foo', '\"bar baz\"']"
+
 def test_arglist_decorator_twice(argparse_app):
     out = run_cmd(argparse_app, 'arglisttwice "we  should" get these')
     assert out[0] == 'we  should get these'