| OLD | NEW |
| 1 # Copyright 2013 The Chromium Authors. All rights reserved. | 1 # Copyright 2013 The Chromium Authors. All rights reserved. |
| 2 # Use of this source code is governed under the Apache License, Version 2.0 that | 2 # Use of this source code is governed by a BSD-style license that can be |
| 3 # can be found in the LICENSE file. | 3 # found in the LICENSE file. |
| 4 | 4 |
| 5 """Manages subcommands in a script. | 5 """Manages subcommands in a script. |
| 6 | 6 |
| 7 Each subcommand should look like this: | 7 Each subcommand should look like this: |
| 8 @usage('[pet name]') | 8 @usage('[pet name]') |
| 9 def CMDpet(parser, args): | 9 def CMDpet(parser, args): |
| 10 '''Prints a pet. | 10 '''Prints a pet. |
| 11 | 11 |
| 12 Many people likes pet. This command prints a pet for your pleasure. | 12 Many people likes pet. This command prints a pet for your pleasure. |
| 13 ''' | 13 ''' |
| (...skipping 16 matching lines...) Expand all Loading... |
| 30 - Every function in the specified module with a name starting with 'CMD' will | 30 - Every function in the specified module with a name starting with 'CMD' will |
| 31 be a subcommand. | 31 be a subcommand. |
| 32 - The module's docstring will be used in the default 'help' page. | 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. | 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. | 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.: | 35 - If a command is an alias to another one, it won't be documented. E.g.: |
| 36 CMDoldname = CMDnewcmd | 36 CMDoldname = CMDnewcmd |
| 37 will result in oldname not being documented but supported and redirecting to | 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 | 38 newcmd. Make it a real function that calls the old function if you want it |
| 39 to be documented. | 39 to be documented. |
| 40 - CMDfoo_bar will be command 'foo-bar'. |
| 40 """ | 41 """ |
| 41 | 42 |
| 42 import difflib | 43 import difflib |
| 43 import sys | 44 import sys |
| 44 import textwrap | 45 import textwrap |
| 45 | 46 |
| 46 | 47 |
| 47 def usage(more): | 48 def usage(more): |
| 48 """Adds a 'usage_more' property to a CMD function.""" | 49 """Adds a 'usage_more' property to a CMD function.""" |
| 49 def hook(fn): | 50 def hook(fn): |
| (...skipping 25 matching lines...) Expand all Loading... |
| 75 | 76 |
| 76 | 77 |
| 77 def _get_color_module(): | 78 def _get_color_module(): |
| 78 """Returns the colorama module if available. | 79 """Returns the colorama module if available. |
| 79 | 80 |
| 80 If so, assumes colors are supported and return the module handle. | 81 If so, assumes colors are supported and return the module handle. |
| 81 """ | 82 """ |
| 82 return sys.modules.get('colorama') or sys.modules.get('third_party.colorama') | 83 return sys.modules.get('colorama') or sys.modules.get('third_party.colorama') |
| 83 | 84 |
| 84 | 85 |
| 86 def _function_to_name(name): |
| 87 """Returns the name of a CMD function.""" |
| 88 return name[3:].replace('_', '-') |
| 89 |
| 90 |
| 85 class CommandDispatcher(object): | 91 class CommandDispatcher(object): |
| 86 def __init__(self, module): | 92 def __init__(self, module): |
| 87 """module is the name of the main python module where to look for commands. | 93 """module is the name of the main python module where to look for commands. |
| 88 | 94 |
| 89 The python builtin variable __name__ MUST be used for |module|. If the | 95 The python builtin variable __name__ MUST be used for |module|. If the |
| 90 script is executed in the form 'python script.py', __name__ == '__main__' | 96 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 | 97 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 | 98 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. | 99 itself with 'script'. __name__ always match the right value. |
| 94 """ | 100 """ |
| 95 self.module = sys.modules[module] | 101 self.module = sys.modules[module] |
| 96 | 102 |
| 97 def enumerate_commands(self): | 103 def enumerate_commands(self): |
| 98 """Returns a dict of command and their handling function. | 104 """Returns a dict of command and their handling function. |
| 99 | 105 |
| 100 The commands must be in the '__main__' modules. To import a command from a | 106 The commands must be in the '__main__' modules. To import a command from a |
| 101 submodule, use: | 107 submodule, use: |
| 102 from mysubcommand import CMDfoo | 108 from mysubcommand import CMDfoo |
| 103 | 109 |
| 104 Automatically adds 'help' if not already defined. | 110 Automatically adds 'help' if not already defined. |
| 105 | 111 |
| 112 Normalizes '_' in the commands to '-'. |
| 113 |
| 106 A command can be effectively disabled by defining a global variable to None, | 114 A command can be effectively disabled by defining a global variable to None, |
| 107 e.g.: | 115 e.g.: |
| 108 CMDhelp = None | 116 CMDhelp = None |
| 109 """ | 117 """ |
| 110 cmds = dict( | 118 cmds = dict( |
| 111 (fn[3:], getattr(self.module, fn)) | 119 (_function_to_name(name), getattr(self.module, name)) |
| 112 for fn in dir(self.module) if fn.startswith('CMD')) | 120 for name in dir(self.module) if name.startswith('CMD')) |
| 113 cmds.setdefault('help', CMDhelp) | 121 cmds.setdefault('help', CMDhelp) |
| 114 return cmds | 122 return cmds |
| 115 | 123 |
| 116 def find_nearest_command(self, name): | 124 def find_nearest_command(self, name_asked): |
| 117 """Retrieves the function to handle a command. | 125 """Retrieves the function to handle a command as supplied by the user. |
| 118 | 126 |
| 119 It automatically tries to guess the intended command by handling typos or | 127 It automatically tries to guess the _intended command_ by handling typos |
| 120 incomplete names. | 128 and/or incomplete names. |
| 121 """ | 129 """ |
| 122 commands = self.enumerate_commands() | 130 commands = self.enumerate_commands() |
| 123 if name in commands: | 131 if name_asked in commands: |
| 124 return commands[name] | 132 return commands[name_asked] |
| 125 | 133 |
| 126 # An exact match was not found. Try to be smart and look if there's | 134 # An exact match was not found. Try to be smart and look if there's |
| 127 # something similar. | 135 # something similar. |
| 128 commands_with_prefix = [c for c in commands if c.startswith(name)] | 136 commands_with_prefix = [c for c in commands if c.startswith(name_asked)] |
| 129 if len(commands_with_prefix) == 1: | 137 if len(commands_with_prefix) == 1: |
| 130 return commands[commands_with_prefix[0]] | 138 return commands[commands_with_prefix[0]] |
| 131 | 139 |
| 132 # A #closeenough approximation of levenshtein distance. | 140 # A #closeenough approximation of levenshtein distance. |
| 133 def close_enough(a, b): | 141 def close_enough(a, b): |
| 134 return difflib.SequenceMatcher(a=a, b=b).ratio() | 142 return difflib.SequenceMatcher(a=a, b=b).ratio() |
| 135 | 143 |
| 136 hamming_commands = sorted( | 144 hamming_commands = sorted( |
| 137 ((close_enough(c, name), c) for c in commands), | 145 ((close_enough(c, name_asked), c) for c in commands), |
| 138 reverse=True) | 146 reverse=True) |
| 139 if (hamming_commands[0][0] - hamming_commands[1][0]) < 0.3: | 147 if (hamming_commands[0][0] - hamming_commands[1][0]) < 0.3: |
| 140 # Too ambiguous. | 148 # Too ambiguous. |
| 141 return | 149 return |
| 142 | 150 |
| 143 if hamming_commands[0][0] < 0.8: | 151 if hamming_commands[0][0] < 0.8: |
| 144 # Not similar enough. Don't be a fool and run a random command. | 152 # Not similar enough. Don't be a fool and run a random command. |
| 145 return | 153 return |
| 146 | 154 |
| 147 return commands[hamming_commands[0][1]] | 155 return commands[hamming_commands[0][1]] |
| 148 | 156 |
| 149 def _gen_commands_list(self): | 157 def _gen_commands_list(self): |
| 150 """Generates the short list of supported commands.""" | 158 """Generates the short list of supported commands.""" |
| 151 commands = self.enumerate_commands() | 159 commands = self.enumerate_commands() |
| 152 docs = sorted( | 160 docs = sorted( |
| 153 (name, self._create_command_summary(name, handler)) | 161 (cmd_name, self._create_command_summary(cmd_name, handler)) |
| 154 for name, handler in commands.iteritems()) | 162 for cmd_name, handler in commands.iteritems()) |
| 155 # Skip commands without a docstring. | 163 # Skip commands without a docstring. |
| 156 docs = [i for i in docs if i[1]] | 164 docs = [i for i in docs if i[1]] |
| 157 # Then calculate maximum length for alignment: | 165 # Then calculate maximum length for alignment: |
| 158 length = max(len(c) for c in commands) | 166 length = max(len(c) for c in commands) |
| 159 | 167 |
| 160 # Look if color is supported. | 168 # Look if color is supported. |
| 161 colors = _get_color_module() | 169 colors = _get_color_module() |
| 162 green = reset = '' | 170 green = reset = '' |
| 163 if colors: | 171 if colors: |
| 164 green = colors.Fore.GREEN | 172 green = colors.Fore.GREEN |
| 165 reset = colors.Fore.RESET | 173 reset = colors.Fore.RESET |
| 166 return ( | 174 return ( |
| 167 'Commands are:\n' + | 175 'Commands are:\n' + |
| 168 ''.join( | 176 ''.join( |
| 169 ' %s%-*s%s %s\n' % (green, length, name, reset, doc) | 177 ' %s%-*s%s %s\n' % (green, length, cmd_name, reset, doc) |
| 170 for name, doc in docs)) | 178 for cmd_name, doc in docs)) |
| 171 | 179 |
| 172 def _add_command_usage(self, parser, command): | 180 def _add_command_usage(self, parser, command): |
| 173 """Modifies an OptionParser object with the function's documentation.""" | 181 """Modifies an OptionParser object with the function's documentation.""" |
| 174 name = command.__name__[3:] | 182 cmd_name = _function_to_name(command.__name__) |
| 175 if name == 'help': | 183 if cmd_name == 'help': |
| 176 name = '<command>' | 184 cmd_name = '<command>' |
| 177 # Use the module's docstring as the description for the 'help' command if | 185 # Use the module's docstring as the description for the 'help' command if |
| 178 # available. | 186 # available. |
| 179 parser.description = (self.module.__doc__ or '').rstrip() | 187 parser.description = (self.module.__doc__ or '').rstrip() |
| 180 if parser.description: | 188 if parser.description: |
| 181 parser.description += '\n\n' | 189 parser.description += '\n\n' |
| 182 parser.description += self._gen_commands_list() | 190 parser.description += self._gen_commands_list() |
| 183 # Do not touch epilog. | 191 # Do not touch epilog. |
| 184 else: | 192 else: |
| 185 # Use the command's docstring if available. For commands, unlike module | 193 # Use the command's docstring if available. For commands, unlike module |
| 186 # docstring, realign. | 194 # docstring, realign. |
| 187 lines = (command.__doc__ or '').rstrip().splitlines() | 195 lines = (command.__doc__ or '').rstrip().splitlines() |
| 188 if lines[:1]: | 196 if lines[:1]: |
| 189 rest = textwrap.dedent('\n'.join(lines[1:])) | 197 rest = textwrap.dedent('\n'.join(lines[1:])) |
| 190 parser.description = '\n'.join((lines[0], rest)) | 198 parser.description = '\n'.join((lines[0], rest)) |
| 191 else: | 199 else: |
| 192 parser.description = lines[0] | 200 parser.description = lines[0] if lines else '' |
| 193 if parser.description: | 201 if parser.description: |
| 194 parser.description += '\n' | 202 parser.description += '\n' |
| 195 parser.epilog = getattr(command, 'epilog', None) | 203 parser.epilog = getattr(command, 'epilog', None) |
| 196 if parser.epilog: | 204 if parser.epilog: |
| 197 parser.epilog = '\n' + parser.epilog.strip() + '\n' | 205 parser.epilog = '\n' + parser.epilog.strip() + '\n' |
| 198 | 206 |
| 199 more = getattr(command, 'usage_more', '') | 207 more = getattr(command, 'usage_more', '') |
| 200 parser.set_usage( | 208 extra = '' if not more else ' ' + more |
| 201 'usage: %%prog %s [options]%s' % (name, '' if not more else ' ' + more)) | 209 parser.set_usage('usage: %%prog %s [options]%s' % (cmd_name, extra)) |
| 202 | 210 |
| 203 @staticmethod | 211 @staticmethod |
| 204 def _create_command_summary(name, command): | 212 def _create_command_summary(cmd_name, command): |
| 205 """Creates a oneline summary from the command's docstring.""" | 213 """Creates a oneliner summary from the command's docstring.""" |
| 206 if name != command.__name__[3:]: | 214 if cmd_name != _function_to_name(command.__name__): |
| 207 # Skip aliases. | 215 # Skip aliases. For example using at module level: |
| 216 # CMDfoo = CMDbar |
| 208 return '' | 217 return '' |
| 209 doc = command.__doc__ or '' | 218 doc = command.__doc__ or '' |
| 210 line = doc.split('\n', 1)[0].rstrip('.') | 219 line = doc.split('\n', 1)[0].rstrip('.') |
| 211 if not line: | 220 if not line: |
| 212 return line | 221 return line |
| 213 return (line[0].lower() + line[1:]).strip() | 222 return (line[0].lower() + line[1:]).strip() |
| 214 | 223 |
| 215 def execute(self, parser, args): | 224 def execute(self, parser, args): |
| 216 """Dispatches execution to the right command. | 225 """Dispatches execution to the right command. |
| 217 | 226 |
| (...skipping 24 matching lines...) Expand all Loading... |
| 242 return command(parser, args[1:]) | 251 return command(parser, args[1:]) |
| 243 | 252 |
| 244 cmdhelp = self.enumerate_commands().get('help') | 253 cmdhelp = self.enumerate_commands().get('help') |
| 245 if cmdhelp: | 254 if cmdhelp: |
| 246 # Not a known command. Default to help. | 255 # Not a known command. Default to help. |
| 247 self._add_command_usage(parser, cmdhelp) | 256 self._add_command_usage(parser, cmdhelp) |
| 248 return cmdhelp(parser, args) | 257 return cmdhelp(parser, args) |
| 249 | 258 |
| 250 # Nothing can be done. | 259 # Nothing can be done. |
| 251 return 2 | 260 return 2 |
| OLD | NEW |