| 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 |
|---|
Chromium Code Reviews
