view profilemanager/command.py @ 66:67dc95a355cc

finish basic restore logic
author Jeff Hammel <jhammel@mozilla.com>
date Fri, 07 May 2010 14:56:59 -0700
parents 18f16bd1ba6b
children 145e111903d2
line wrap: on
line source

"""
a command-line interface to the command line, a la pythonpaste
"""

import inspect
import sys
from optparse import OptionParser
from pprint import pprint

if 'commands' not in globals():
    commands = {}
    
def command(function):
    # XXX should get bound/unbound state from function (how?)
    global commands
    name = function.func_name
    doc = inspect.cleandoc(function.__doc__)
    argspec = inspect.getargspec(function)
    defaults = argspec.defaults
    if defaults:
        args = argspec.args[1:-len(defaults)]
        optional = dict(zip(argspec.args[-len(defaults):], defaults))
    else:
        args = argspec.args[1:]
        optional = None
    commands[name] = { 'doc': doc,
                       'args': args, 
                       'optional': optional,
                       'varargs': argspec.varargs
                       }
    return function

def commandargs2str(command):
    if isinstance(command, basestring):
        command = commands[command]
    retval = []
    retval.extend(['<%s>' % arg for arg in command['args']])
    varargs = command['varargs']
    if varargs:
        retval.append('<%s> [%s] [...]' % (varargs, varargs))
    if command['optional']:
        retval.append('[options]')
    return ' '.join(retval)


def doc2arghelp(docstring, decoration='-', delimeter=':'):
    """
    Parse a docstring and get at the section describing arguments
    - decoration: decoration character
    - delimeter: delimter character

    Yields a tuple of the stripped docstring and the arguments help
    dictionary
    """
    lines = [ i.strip() for i in docstring.split('\n') ]
    argdict = {}
    doc = []
    option = None
    for line in lines:
        if not line and option: # blank lines terminate
            break
        if line.startswith(decoration) and delimeter in line:
            name, description = line.split(delimeter, 1)
            name = name.lstrip(decoration).strip()
            description = description.strip()
            argdict[name] = [ description ]
            option = name
        else:
            if option:
                argdict[name].append(line)
            else:
                doc.append(line)
    argdict = dict([(key, ' '.join(value))
                    for key, value in argdict.items()])
    return ('\n'.join(doc), argdict)

def command2parser(command):
    doc, argdict = doc2arghelp(commands[command]['doc'])
    parser = OptionParser('%%prog %s %s' % (command, commandargs2str(command)),
                          description=doc, add_help_option=False)
    if commands[command]['optional']:
        for key, value in commands[command]['optional'].items():
            help = argdict.get(key)
            if value is True:
                parser.add_option('--no-%s' % key, dest=key,
                                  action='store_false', default=True,
                                  help=help)
            elif value is False:
                parser.add_option('--%s' % key, action='store_true',
                                  default=False, help=help)
            else:
                parser.add_option('--%s' % key, help=help)
                                  
    return parser

class CommandParser(OptionParser):
    def __init__(self, commands, description=None, setup=None):
        usage = '%prog [options] command [command-options]'
        OptionParser.__init__(self, usage=usage, description=description)
        for _command in commands:
            command(_command)
        self.disable_interspersed_args()
        self.setup = setup

    def print_help(self):
        OptionParser.print_help(self)
        # short descriptions for commands
        command_descriptions = [dict(name=i,
                                     description=commands[i]['doc'].strip().split('\n',1)[0])
                                for i in sorted(commands.keys())]
        max_len = max([len(i['name']) for i in command_descriptions])
        description = "Commands: \n%s" % ('\n'.join(['  %s%s  %s' % (description['name'], ' ' * (max_len - len(description['name'])), description['description'])
                                                     for description in command_descriptions]))

        print
        print description

    def parse(self, args=sys.argv[1:]):
        """global parse step"""
        
        self.options, args = self.parse_args(args)

        # help/sanity check -- should probably be separated
        if not len(args):
            self.print_help()
            sys.exit(0)
        if args[0] == 'help':
            if len(args) == 2:
                if args[1] in commands:
                    name = args[1]
                    commandparser = command2parser(name)
                    commandparser.print_help()
                else:
                    self.error("No command '%s'" % args[1])
            else:
                self.print_help()
            sys.exit(0)
        command = args[0]
        if command not in commands:
            self.error("No command '%s'" % command)
        return command, args[1:]

    def invoke(self, args=sys.argv[1:]):
        """
        invoke
        """

        # parse
        name, args = self.parse(args)

        # setup
        _object = self.setup(self, self.options)

        # command specific args
        command = commands[name]
        commandparser = command2parser(name)
        command_options, command_args = commandparser.parse_args(args)
        if len(command_args) < len(command['args']):
            commandparser.error("Not enough arguments given")
        if len(command_args) != len(command['args']) and not command['varargs']:
            commandparser.error("Too many arguments given")

        # invoke the command
        retval = getattr(_object, name)(*command_args, **command_options.__dict__)

        # print the output
        if retval is None:
            pass
        elif isinstance(retval, basestring):
            print retval
        elif isinstance(retval, dict):
            for key in sorted(retval.keys()):
                print '%s: %s' % (key, retval[key])
        elif hasattr(retval, '__iter__'):

            # hack since python doesn't have ordered dicts
            if not [ val for val in retval
                     if not(isinstance(val, tuple) and len(val) == 2) ]:
                for val in retval:
                    print '%s: %s' % (val[0], val[1])

            else:
                for val in retval:
                    print val
        else:
            pprint(retval)

        # return the value
        return retval