1 files changed, 97 insertions, 7 deletions
diff --git a/rpkid/rpki/cli.py b/rpkid/rpki/cli.py
index ad485b6f..7d844509 100644
--- a/rpkid/rpki/cli.py
+++ b/rpkid/rpki/cli.py
@@ -23,6 +23,7 @@ Utilities for writing command line tools.
 import cmd
 import glob
 import os.path
+import argparse
 import traceback
 
 try:
@@ -34,6 +35,13 @@ except ImportError:
 class BadCommandSyntax(Exception):
   "Bad command line syntax."
 
+class ExitArgparse(Exception):
+  "Exit method from ArgumentParser."
+
+  def __init__(self, message = None, status = 0):
+    self.message = message
+    self.status = status
+
 class Cmd(cmd.Cmd):
   """
   Customized subclass of Python cmd module.
@@ -61,11 +69,17 @@ class Cmd(cmd.Cmd):
     Wrap error handling around cmd.Cmd.onecmd().  Might want to do
     something kinder than showing a traceback, eventually.
     """
+
     self.last_command_failed = False
     try:
       return cmd.Cmd.onecmd(self, line)
     except SystemExit:
       raise
+    except ExitArgparse, e:
+      if e.message is not None:
+        print e.message
+      self.last_command_failed = e.status != 0
+      return False
     except BadCommandSyntax, e:
       print e
     except Exception:
@@ -74,9 +88,6 @@ class Cmd(cmd.Cmd):
     return False
 
   def do_EOF(self, arg):
-    """
-    Exit program.
-    """
     if self.EOF_exits_command_loop and self.prompt:
       print
     return self.EOF_exits_command_loop
@@ -85,6 +96,7 @@ class Cmd(cmd.Cmd):
     """
     Exit program.
     """
+
     return True
 
   do_quit = do_exit
@@ -95,6 +107,7 @@ class Cmd(cmd.Cmd):
     command, which I find to be violation of the principal of least
     astonishment, so my preference is that an empty line does nothing.
     """
+
     if self.emptyline_repeats_last_command:
       cmd.Cmd.emptyline(self)
 
@@ -104,6 +117,7 @@ class Cmd(cmd.Cmd):
     the normal (bash-like) behavior when one hits the completion key
     and there's only one match.
     """
+
     result = glob.glob(text + "*")
     if len(result) == 1:
       path = result.pop()
@@ -119,27 +133,30 @@ class Cmd(cmd.Cmd):
     consider the normal (bash-like) behavior when one hits the
     completion key and there's only one match.
     """
-    result = set(cmd.Cmd.completenames(self, text, *ignored))
+
+    result = cmd.Cmd.completenames(self, text, *ignored)
     if len(result) == 1:
-      result.add(result.pop() + " ")
-    return list(result)
+      result[0] += " "
+    return result
 
   def help_help(self):
     """
     Type "help [topic]" for help on a command,
     or just "help" for a list of commands.
     """
+
     self.stdout.write(self.help_help.__doc__ + "\n")
 
   def complete_help(self, *args):
     """
     Better completion function for help command arguments.
     """
+
     text = args[0]
     names = self.get_names()
     result = []
     for prefix in ("do_", "help_"):
-      result.extend(s[len(prefix):] for s in names if s.startswith(prefix + text))
+      result.extend(s[len(prefix):] for s in names if s.startswith(prefix + text) and s != "do_EOF")
     return result
 
   if have_readline:
@@ -149,6 +166,7 @@ class Cmd(cmd.Cmd):
       Better command loop, with history file and tweaked readline
       completion delimiters.
       """
+
       old_completer_delims = readline.get_completer_delims()
       if self.histfile is not None:
         try:
@@ -173,6 +191,7 @@ def yes_or_no(prompt, default = None, require_full_word = False):
   """
   Ask a yes-or-no question.
   """
+
   prompt = prompt.rstrip() + _yes_or_no_prompts[default]
   while True:
     answer = raw_input(prompt).strip().lower()
@@ -189,3 +208,74 @@ _yes_or_no_prompts = {
   False : ' ("yes" or "no" ["no"]) ',
   None  : ' ("yes" or "no") ' }
 
+
+class NonExitingArgumentParser(argparse.ArgumentParser):
+  """
+  ArgumentParser tweaked to throw ExitArgparse exception
+  rather than using sys.exit(), for use with command loop.
+  """
+
+  def exit(self, status = 0, message = None):
+    raise ExitArgparse(status = status, message = message)
+
+
+def parsecmd(subparsers, *arg_clauses):
+  """
+  Decorator to combine the argparse and cmd modules.
+
+  subparsers is an instance of argparse.ArgumentParser (or subclass) which was
+  returned by calling the .add_subparsers() method on an ArgumentParser instance
+  intended to handle parsing for the entire program on the command line.
+
+  arg_clauses is a series of defarg() invocations defining arguments to be parsed
+  by the argparse code. 
+
+  The decorator will use arg_clauses to construct two separate argparse parser
+  instances: one will be attached to the global parser as a subparser, the
+  other will be used to parse arguments for this command when invoked by cmd.
+
+  The decorator will replace the original do_whatever method with a wrapped version
+  which uses the local argparse instance to parse the single string supplied by
+  the cmd module.
+
+  The intent is that, from the command's point of view, all of this should work
+  pretty much the same way regardless of whether the command was invoked from
+  the global command line or from within the cmd command loop.  Either way,
+  the command method should get an argparse.Namespace object.
+
+  In theory, we could generate a completion handler from the argparse definitions,
+  much as the separate argcomplete package does.  In practice this is a lot of
+  work and I'm not ready to get into that just yet.
+  """
+
+  # We probably want to use a customized subclass of ArgumentParser to
+  # avoid the exit-on-error behavior.  We may also want to use something
+  # fancier than plain str.split() to split arguments.  Ignore all this
+  # for the moment, yak shaving.
+
+  def decorate(func):
+    assert func.__name__.startswith("do_")
+    parser = NonExitingArgumentParser(description = func.__doc__,
+                                      prog = func.__name__[3:],
+                                      add_help = False)
+    subparser = subparsers.add_parser(func.__name__[3:],
+                                      description = func.__doc__,
+                                      help = func.__doc__.lstrip().partition("\n")[0])
+    for positional, keywords in arg_clauses:
+      parser.add_argument(*positional, **keywords)
+      subparser.add_argument(*positional, **keywords)
+    subparser.set_defaults(func = func)
+    def wrapped(self, arg):
+      return func(self, parser.parse_args(arg.split()))
+    wrapped.argparser = parser
+    wrapped.__doc__ = func.__doc__
+    return wrapped
+  return decorate
+
+def cmdarg(*positional, **keywords):
+  """
+  Syntactic sugar to let us use keyword arguments normally when constructing
+  arguments for deferred calls to argparse.ArgumentParser.add_argument().
+  """
+
+  return positional, keywords