blob: a9ebae098d49babcd3693ea6a1a570b75cc10e50 [file] [log] [blame]
# Copyright 2013 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Manages subcommands in a script.
Each subcommand should look like this:
@usage('[pet name]')
def CMDpet(parser, args):
'''Prints a pet.
Many people likes pet. This command prints a pet for your pleasure.
parser.add_option('--color', help='color of your pet')
options, args = parser.parse_args(args)
if len(args) != 1:
parser.error('A pet name is required')
pet = args[0]
if options.color:
print('Nice %s %d' % (options.color, pet))
print('Nice %s' % pet)
return 0
- usage decorator alters the 'usage: %prog' line in the command's help.
- docstring is used to both short help line and long help line.
- parser can be augmented with arguments.
- return the exit code.
- Every function in the specified module with a name starting with 'CMD' will
be a subcommand.
- The module's docstring will be used in the default 'help' page.
- If a command has no docstring, it will not be listed in the 'help' page.
Useful to keep compatibility commands around or aliases.
- If a command is an alias to another one, it won't be documented. E.g.:
CMDoldname = CMDnewcmd
will result in oldname not being documented but supported and redirecting to
newcmd. Make it a real function that calls the old function if you want it
to be documented.
- CMDfoo_bar will be command 'foo-bar'.
import difflib
import sys
import textwrap
def usage(more):
"""Adds a 'usage_more' property to a CMD function."""
def hook(fn):
fn.usage_more = more
return fn
return hook
def epilog(text):
"""Adds an 'epilog' property to a CMD function.
It will be shown in the epilog. Usually useful for examples.
def hook(fn):
fn.epilog = text
return fn
return hook
def CMDhelp(parser, args):
"""Prints list of commands or help for a specific command."""
# This is the default help implementation. It can be disabled or overriden if
# wanted.
if not any(i in ('-h', '--help') for i in args):
args = args + ['--help']
_, args = parser.parse_args(args)
# Never gets there.
assert False
def _get_color_module():
"""Returns the colorama module if available.
If so, assumes colors are supported and return the module handle.
return sys.modules.get('colorama') or sys.modules.get('third_party.colorama')
def _function_to_name(name):
"""Returns the name of a CMD function."""
return name[3:].replace('_', '-')
class CommandDispatcher(object):
def __init__(self, module):
"""module is the name of the main python module where to look for commands.
The python builtin variable __name__ MUST be used for |module|. If the
script is executed in the form 'python', __name__ == '__main__'
and sys.modules['script'] doesn't exist. On the other hand if it is unit
tested, __main__ will be the unit test's module so it has to reference to
itself with 'script'. __name__ always match the right value.
self.module = sys.modules[module]
def enumerate_commands(self):
"""Returns a dict of command and their handling function.
The commands must be in the '__main__' modules. To import a command from a
submodule, use:
from mysubcommand import CMDfoo
Automatically adds 'help' if not already defined.
Normalizes '_' in the commands to '-'.
A command can be effectively disabled by defining a global variable to None,
CMDhelp = None
cmds = dict(
(_function_to_name(name), getattr(self.module, name))
for name in dir(self.module) if name.startswith('CMD'))
cmds.setdefault('help', CMDhelp)
return cmds
def find_nearest_command(self, name_asked):
"""Retrieves the function to handle a command as supplied by the user.
It automatically tries to guess the _intended command_ by handling typos
and/or incomplete names.
commands = self.enumerate_commands()
name_to_dash = name_asked.replace('_', '-')
if name_to_dash in commands:
return commands[name_to_dash]
# An exact match was not found. Try to be smart and look if there's
# something similar.
commands_with_prefix = [c for c in commands if c.startswith(name_asked)]
if len(commands_with_prefix) == 1:
return commands[commands_with_prefix[0]]
# A #closeenough approximation of levenshtein distance.
def close_enough(a, b):
return difflib.SequenceMatcher(a=a, b=b).ratio()
hamming_commands = sorted(
((close_enough(c, name_asked), c) for c in commands),
if (hamming_commands[0][0] - hamming_commands[1][0]) < 0.3:
# Too ambiguous.
if hamming_commands[0][0] < 0.8:
# Not similar enough. Don't be a fool and run a random command.
return commands[hamming_commands[0][1]]
def _gen_commands_list(self):
"""Generates the short list of supported commands."""
commands = self.enumerate_commands()
docs = sorted(
(cmd_name, self._create_command_summary(cmd_name, handler))
for cmd_name, handler in commands.iteritems())
# Skip commands without a docstring.
docs = [i for i in docs if i[1]]
# Then calculate maximum length for alignment:
length = max(len(c) for c in commands)
# Look if color is supported.
colors = _get_color_module()
green = reset = ''
if colors:
green = colors.Fore.GREEN
reset = colors.Fore.RESET
return (
'Commands are:\n' +
' %s%-*s%s %s\n' % (green, length, cmd_name, reset, doc)
for cmd_name, doc in docs))
def _add_command_usage(self, parser, command):
"""Modifies an OptionParser object with the function's documentation."""
cmd_name = _function_to_name(command.__name__)
if cmd_name == 'help':
cmd_name = '<command>'
# Use the module's docstring as the description for the 'help' command if
# available.
parser.description = (self.module.__doc__ or '').rstrip()
if parser.description:
parser.description += '\n\n'
parser.description += self._gen_commands_list()
# Do not touch epilog.
# Use the command's docstring if available. For commands, unlike module
# docstring, realign.
lines = (command.__doc__ or '').rstrip().splitlines()
if lines[:1]:
rest = textwrap.dedent('\n'.join(lines[1:]))
parser.description = '\n'.join((lines[0], rest))
parser.description = lines[0] if lines else ''
if parser.description:
parser.description += '\n'
parser.epilog = getattr(command, 'epilog', None)
if parser.epilog:
parser.epilog = '\n' + parser.epilog.strip() + '\n'
more = getattr(command, 'usage_more', '')
extra = '' if not more else ' ' + more
parser.set_usage('usage: %%prog %s [options]%s' % (cmd_name, extra))
def _create_command_summary(cmd_name, command):
"""Creates a oneliner summary from the command's docstring."""
if cmd_name != _function_to_name(command.__name__):
# Skip aliases. For example using at module level:
# CMDfoo = CMDbar
return ''
doc = command.__doc__ or ''
line = doc.split('\n', 1)[0].rstrip('.')
if not line:
return line
return (line[0].lower() + line[1:]).strip()
def execute(self, parser, args):
"""Dispatches execution to the right command.
Fallbacks to 'help' if not disabled.
# Unconditionally disable format_description() and format_epilog().
# Technically, a formatter should be used but it's not worth (yet) the
# trouble.
parser.format_description = lambda _: parser.description or ''
parser.format_epilog = lambda _: parser.epilog or ''
if args:
if args[0] in ('-h', '--help') and len(args) > 1:
# Inverse the argument order so 'tool --help cmd' is rewritten to
# 'tool cmd --help'.
args = [args[1], args[0]] + args[2:]
command = self.find_nearest_command(args[0])
if command:
if command.__name__ == 'CMDhelp' and len(args) > 1:
# Inverse the arguments order so 'tool help cmd' is rewritten to
# 'tool cmd --help'. Do it here since we want 'tool hel cmd' to work
# too.
args = [args[1], '--help'] + args[2:]
command = self.find_nearest_command(args[0]) or command
# "fix" the usage and the description now that we know the subcommand.
self._add_command_usage(parser, command)
return command(parser, args[1:])
cmdhelp = self.enumerate_commands().get('help')
if cmdhelp:
# Not a known command. Default to help.
self._add_command_usage(parser, cmdhelp)
return cmdhelp(parser, args)
# Nothing can be done.
return 2