OLD | NEW |
(Empty) | |
| 1 # Copyright 2013 The Chromium Authors. All rights reserved. |
| 2 # Use of this source code is governed by a BSD-style license that can be |
| 3 # found in the LICENSE file. |
| 4 |
| 5 """Manages subcommands in a script. |
| 6 |
| 7 Each subcommand should look like this: |
| 8 @usage('[pet name]') |
| 9 def CMDpet(parser, args): |
| 10 '''Prints a pet. |
| 11 |
| 12 Many people likes pet. This command prints a pet for your pleasure. |
| 13 ''' |
| 14 parser.add_option('--color', help='color of your pet') |
| 15 options, args = parser.parse_args(args) |
| 16 if len(args) != 1: |
| 17 parser.error('A pet name is required') |
| 18 pet = args[0] |
| 19 if options.color: |
| 20 print('Nice %s %d' % (options.color, pet)) |
| 21 else: |
| 22 print('Nice %s' % pet) |
| 23 return 0 |
| 24 |
| 25 Explanation: |
| 26 - usage decorator alters the 'usage: %prog' line in the command's help. |
| 27 - docstring is used to both short help line and long help line. |
| 28 - parser can be augmented with arguments. |
| 29 - return the exit code. |
| 30 - Every function in the specified module with a name starting with 'CMD' will |
| 31 be a subcommand. |
| 32 - The module's docstring will be used in the default 'help' page. |
| 33 - If a command has no docstring, it will not be listed in the 'help' page. |
| 34 Useful to keep compatibility commands around or aliases. |
| 35 - If a command is an alias to another one, it won't be documented. E.g.: |
| 36 CMDoldname = CMDnewcmd |
| 37 will result in oldname not being documented but supported and redirecting to |
| 38 newcmd. Make it a real function that calls the old function if you want it |
| 39 to be documented. |
| 40 """ |
| 41 |
| 42 import difflib |
| 43 import sys |
| 44 import textwrap |
| 45 |
| 46 |
| 47 def usage(more): |
| 48 """Adds a 'usage_more' property to a CMD function.""" |
| 49 def hook(fn): |
| 50 fn.usage_more = more |
| 51 return fn |
| 52 return hook |
| 53 |
| 54 |
| 55 def epilog(text): |
| 56 """Adds an 'epilog' property to a CMD function. |
| 57 |
| 58 It will be shown in the epilog. Usually useful for examples. |
| 59 """ |
| 60 def hook(fn): |
| 61 fn.epilog = text |
| 62 return fn |
| 63 return hook |
| 64 |
| 65 |
| 66 def CMDhelp(parser, args): |
| 67 """Prints list of commands or help for a specific command.""" |
| 68 # This is the default help implementation. It can be disabled or overriden if |
| 69 # wanted. |
| 70 if not any(i in ('-h', '--help') for i in args): |
| 71 args = args + ['--help'] |
| 72 _, args = parser.parse_args(args) |
| 73 # Never gets there. |
| 74 assert False |
| 75 |
| 76 |
| 77 def _get_color_module(): |
| 78 """Returns the colorama module if available. |
| 79 |
| 80 If so, assumes colors are supported and return the module handle. |
| 81 """ |
| 82 return sys.modules.get('colorama') or sys.modules.get('third_party.colorama') |
| 83 |
| 84 |
| 85 class CommandDispatcher(object): |
| 86 def __init__(self, module): |
| 87 """module is the name of the main python module where to look for commands. |
| 88 |
| 89 The python builtin variable __name__ MUST be used for |module|. If the |
| 90 script is executed in the form 'python script.py', __name__ == '__main__' |
| 91 and sys.modules['script'] doesn't exist. On the other hand if it is unit |
| 92 tested, __main__ will be the unit test's module so it has to reference to |
| 93 itself with 'script'. __name__ always match the right value. |
| 94 """ |
| 95 self.module = sys.modules[module] |
| 96 |
| 97 def enumerate_commands(self): |
| 98 """Returns a dict of command and their handling function. |
| 99 |
| 100 The commands must be in the '__main__' modules. To import a command from a |
| 101 submodule, use: |
| 102 from mysubcommand import CMDfoo |
| 103 |
| 104 Automatically adds 'help' if not already defined. |
| 105 |
| 106 A command can be effectively disabled by defining a global variable to None, |
| 107 e.g.: |
| 108 CMDhelp = None |
| 109 """ |
| 110 cmds = dict( |
| 111 (fn[3:], getattr(self.module, fn)) |
| 112 for fn in dir(self.module) if fn.startswith('CMD')) |
| 113 cmds.setdefault('help', CMDhelp) |
| 114 return cmds |
| 115 |
| 116 def find_nearest_command(self, name): |
| 117 """Retrieves the function to handle a command. |
| 118 |
| 119 It automatically tries to guess the intended command by handling typos or |
| 120 incomplete names. |
| 121 """ |
| 122 commands = self.enumerate_commands() |
| 123 if name in commands: |
| 124 return commands[name] |
| 125 |
| 126 # An exact match was not found. Try to be smart and look if there's |
| 127 # something similar. |
| 128 commands_with_prefix = [c for c in commands if c.startswith(name)] |
| 129 if len(commands_with_prefix) == 1: |
| 130 return commands[commands_with_prefix[0]] |
| 131 |
| 132 # A #closeenough approximation of levenshtein distance. |
| 133 def close_enough(a, b): |
| 134 return difflib.SequenceMatcher(a=a, b=b).ratio() |
| 135 |
| 136 hamming_commands = sorted( |
| 137 ((close_enough(c, name), c) for c in commands), |
| 138 reverse=True) |
| 139 if (hamming_commands[0][0] - hamming_commands[1][0]) < 0.3: |
| 140 # Too ambiguous. |
| 141 return |
| 142 |
| 143 if hamming_commands[0][0] < 0.8: |
| 144 # Not similar enough. Don't be a fool and run a random command. |
| 145 return |
| 146 |
| 147 return commands[hamming_commands[0][1]] |
| 148 |
| 149 def _gen_commands_list(self): |
| 150 """Generates the short list of supported commands.""" |
| 151 commands = self.enumerate_commands() |
| 152 docs = sorted( |
| 153 (name, self._create_command_summary(name, handler)) |
| 154 for name, handler in commands.iteritems()) |
| 155 # Skip commands without a docstring. |
| 156 docs = [i for i in docs if i[1]] |
| 157 # Then calculate maximum length for alignment: |
| 158 length = max(len(c) for c in commands) |
| 159 |
| 160 # Look if color is supported. |
| 161 colors = _get_color_module() |
| 162 green = reset = '' |
| 163 if colors: |
| 164 green = colors.Fore.GREEN |
| 165 reset = colors.Fore.RESET |
| 166 return ( |
| 167 'Commands are:\n' + |
| 168 ''.join( |
| 169 ' %s%-*s%s %s\n' % (green, length, name, reset, doc) |
| 170 for name, doc in docs)) |
| 171 |
| 172 def _add_command_usage(self, parser, command): |
| 173 """Modifies an OptionParser object with the function's documentation.""" |
| 174 name = command.__name__[3:] |
| 175 if name == 'help': |
| 176 name = '<command>' |
| 177 # Use the module's docstring as the description for the 'help' command if |
| 178 # available. |
| 179 parser.description = (self.module.__doc__ or '').rstrip() |
| 180 if parser.description: |
| 181 parser.description += '\n\n' |
| 182 parser.description += self._gen_commands_list() |
| 183 # Do not touch epilog. |
| 184 else: |
| 185 # Use the command's docstring if available. For commands, unlike module |
| 186 # docstring, realign. |
| 187 lines = (command.__doc__ or '').rstrip().splitlines() |
| 188 if lines[:1]: |
| 189 rest = textwrap.dedent('\n'.join(lines[1:])) |
| 190 parser.description = '\n'.join((lines[0], rest)) |
| 191 else: |
| 192 parser.description = lines[0] |
| 193 if parser.description: |
| 194 parser.description += '\n' |
| 195 parser.epilog = getattr(command, 'epilog', None) |
| 196 if parser.epilog: |
| 197 parser.epilog = '\n' + parser.epilog.strip() + '\n' |
| 198 |
| 199 more = getattr(command, 'usage_more', '') |
| 200 parser.set_usage( |
| 201 'usage: %%prog %s [options]%s' % (name, '' if not more else ' ' + more)) |
| 202 |
| 203 @staticmethod |
| 204 def _create_command_summary(name, command): |
| 205 """Creates a oneline summary from the command's docstring.""" |
| 206 if name != command.__name__[3:]: |
| 207 # Skip aliases. |
| 208 return '' |
| 209 doc = command.__doc__ or '' |
| 210 line = doc.split('\n', 1)[0].rstrip('.') |
| 211 if not line: |
| 212 return line |
| 213 return (line[0].lower() + line[1:]).strip() |
| 214 |
| 215 def execute(self, parser, args): |
| 216 """Dispatches execution to the right command. |
| 217 |
| 218 Fallbacks to 'help' if not disabled. |
| 219 """ |
| 220 # Unconditionally disable format_description() and format_epilog(). |
| 221 # Technically, a formatter should be used but it's not worth (yet) the |
| 222 # trouble. |
| 223 parser.format_description = lambda _: parser.description or '' |
| 224 parser.format_epilog = lambda _: parser.epilog or '' |
| 225 |
| 226 if args: |
| 227 if args[0] in ('-h', '--help') and len(args) > 1: |
| 228 # Inverse the argument order so 'tool --help cmd' is rewritten to |
| 229 # 'tool cmd --help'. |
| 230 args = [args[1], args[0]] + args[2:] |
| 231 command = self.find_nearest_command(args[0]) |
| 232 if command: |
| 233 if command.__name__ == 'CMDhelp' and len(args) > 1: |
| 234 # Inverse the arguments order so 'tool help cmd' is rewritten to |
| 235 # 'tool cmd --help'. Do it here since we want 'tool hel cmd' to work |
| 236 # too. |
| 237 args = [args[1], '--help'] + args[2:] |
| 238 command = self.find_nearest_command(args[0]) or command |
| 239 |
| 240 # "fix" the usage and the description now that we know the subcommand. |
| 241 self._add_command_usage(parser, command) |
| 242 return command(parser, args[1:]) |
| 243 |
| 244 cmdhelp = self.enumerate_commands().get('help') |
| 245 if cmdhelp: |
| 246 # Not a known command. Default to help. |
| 247 self._add_command_usage(parser, cmdhelp) |
| 248 return cmdhelp(parser, args) |
| 249 |
| 250 # Nothing can be done. |
| 251 return 2 |
OLD | NEW |