| 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 by a BSD-style license that can be | 2 # Use of this source code is governed by a BSD-style license that can be |
| 3 # 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. |
| (...skipping 19 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 # Implicitly replace foo-bar to foo_bar since foo-bar is not a valid python | |
| 123 # symbol but it's faster to type. | |
| 124 name = name.replace('-', '_') | |
| 125 commands = self.enumerate_commands() | 130 commands = self.enumerate_commands() |
| 126 if name in commands: | 131 if name_asked in commands: |
| 127 return commands[name] | 132 return commands[name_asked] |
| 128 | 133 |
| 129 # 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 |
| 130 # something similar. | 135 # something similar. |
| 131 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)] |
| 132 if len(commands_with_prefix) == 1: | 137 if len(commands_with_prefix) == 1: |
| 133 return commands[commands_with_prefix[0]] | 138 return commands[commands_with_prefix[0]] |
| 134 | 139 |
| 135 # A #closeenough approximation of levenshtein distance. | 140 # A #closeenough approximation of levenshtein distance. |
| 136 def close_enough(a, b): | 141 def close_enough(a, b): |
| 137 return difflib.SequenceMatcher(a=a, b=b).ratio() | 142 return difflib.SequenceMatcher(a=a, b=b).ratio() |
| 138 | 143 |
| 139 hamming_commands = sorted( | 144 hamming_commands = sorted( |
| 140 ((close_enough(c, name), c) for c in commands), | 145 ((close_enough(c, name_asked), c) for c in commands), |
| 141 reverse=True) | 146 reverse=True) |
| 142 if (hamming_commands[0][0] - hamming_commands[1][0]) < 0.3: | 147 if (hamming_commands[0][0] - hamming_commands[1][0]) < 0.3: |
| 143 # Too ambiguous. | 148 # Too ambiguous. |
| 144 return | 149 return |
| 145 | 150 |
| 146 if hamming_commands[0][0] < 0.8: | 151 if hamming_commands[0][0] < 0.8: |
| 147 # 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. |
| 148 return | 153 return |
| 149 | 154 |
| 150 return commands[hamming_commands[0][1]] | 155 return commands[hamming_commands[0][1]] |
| 151 | 156 |
| 152 def _gen_commands_list(self): | 157 def _gen_commands_list(self): |
| 153 """Generates the short list of supported commands.""" | 158 """Generates the short list of supported commands.""" |
| 154 commands = self.enumerate_commands() | 159 commands = self.enumerate_commands() |
| 155 docs = sorted( | 160 docs = sorted( |
| 156 (name, self._create_command_summary(name, handler)) | 161 (cmd_name, self._create_command_summary(cmd_name, handler)) |
| 157 for name, handler in commands.iteritems()) | 162 for cmd_name, handler in commands.iteritems()) |
| 158 # Skip commands without a docstring. | 163 # Skip commands without a docstring. |
| 159 docs = [i for i in docs if i[1]] | 164 docs = [i for i in docs if i[1]] |
| 160 # Then calculate maximum length for alignment: | 165 # Then calculate maximum length for alignment: |
| 161 length = max(len(c) for c in commands) | 166 length = max(len(c) for c in commands) |
| 162 | 167 |
| 163 # Look if color is supported. | 168 # Look if color is supported. |
| 164 colors = _get_color_module() | 169 colors = _get_color_module() |
| 165 green = reset = '' | 170 green = reset = '' |
| 166 if colors: | 171 if colors: |
| 167 green = colors.Fore.GREEN | 172 green = colors.Fore.GREEN |
| 168 reset = colors.Fore.RESET | 173 reset = colors.Fore.RESET |
| 169 return ( | 174 return ( |
| 170 'Commands are:\n' + | 175 'Commands are:\n' + |
| 171 ''.join( | 176 ''.join( |
| 172 ' %s%-*s%s %s\n' % (green, length, name, reset, doc) | 177 ' %s%-*s%s %s\n' % (green, length, cmd_name, reset, doc) |
| 173 for name, doc in docs)) | 178 for cmd_name, doc in docs)) |
| 174 | 179 |
| 175 def _add_command_usage(self, parser, command): | 180 def _add_command_usage(self, parser, command): |
| 176 """Modifies an OptionParser object with the function's documentation.""" | 181 """Modifies an OptionParser object with the function's documentation.""" |
| 177 name = command.__name__[3:] | 182 cmd_name = _function_to_name(command.__name__) |
| 178 if name == 'help': | 183 if cmd_name == 'help': |
| 179 name = '<command>' | 184 cmd_name = '<command>' |
| 180 # 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 |
| 181 # available. | 186 # available. |
| 182 parser.description = (self.module.__doc__ or '').rstrip() | 187 parser.description = (self.module.__doc__ or '').rstrip() |
| 183 if parser.description: | 188 if parser.description: |
| 184 parser.description += '\n\n' | 189 parser.description += '\n\n' |
| 185 parser.description += self._gen_commands_list() | 190 parser.description += self._gen_commands_list() |
| 186 # Do not touch epilog. | 191 # Do not touch epilog. |
| 187 else: | 192 else: |
| 188 # Use the command's docstring if available. For commands, unlike module | 193 # Use the command's docstring if available. For commands, unlike module |
| 189 # docstring, realign. | 194 # docstring, realign. |
| 190 lines = (command.__doc__ or '').rstrip().splitlines() | 195 lines = (command.__doc__ or '').rstrip().splitlines() |
| 191 if lines[:1]: | 196 if lines[:1]: |
| 192 rest = textwrap.dedent('\n'.join(lines[1:])) | 197 rest = textwrap.dedent('\n'.join(lines[1:])) |
| 193 parser.description = '\n'.join((lines[0], rest)) | 198 parser.description = '\n'.join((lines[0], rest)) |
| 194 else: | 199 else: |
| 195 parser.description = lines[0] | 200 parser.description = lines[0] if lines else '' |
| 196 if parser.description: | 201 if parser.description: |
| 197 parser.description += '\n' | 202 parser.description += '\n' |
| 198 parser.epilog = getattr(command, 'epilog', None) | 203 parser.epilog = getattr(command, 'epilog', None) |
| 199 if parser.epilog: | 204 if parser.epilog: |
| 200 parser.epilog = '\n' + parser.epilog.strip() + '\n' | 205 parser.epilog = '\n' + parser.epilog.strip() + '\n' |
| 201 | 206 |
| 202 more = getattr(command, 'usage_more', '') | 207 more = getattr(command, 'usage_more', '') |
| 203 parser.set_usage( | 208 extra = '' if not more else ' ' + more |
| 204 'usage: %%prog %s [options]%s' % (name, '' if not more else ' ' + more)) | 209 parser.set_usage('usage: %%prog %s [options]%s' % (cmd_name, extra)) |
| 205 | 210 |
| 206 @staticmethod | 211 @staticmethod |
| 207 def _create_command_summary(name, command): | 212 def _create_command_summary(cmd_name, command): |
| 208 """Creates a oneline summary from the command's docstring.""" | 213 """Creates a oneliner summary from the command's docstring.""" |
| 209 if name != command.__name__[3:]: | 214 if cmd_name != _function_to_name(command.__name__): |
| 210 # Skip aliases. | 215 # Skip aliases. For example using at module level: |
| 216 # CMDfoo = CMDbar |
| 211 return '' | 217 return '' |
| 212 doc = command.__doc__ or '' | 218 doc = command.__doc__ or '' |
| 213 line = doc.split('\n', 1)[0].rstrip('.') | 219 line = doc.split('\n', 1)[0].rstrip('.') |
| 214 if not line: | 220 if not line: |
| 215 return line | 221 return line |
| 216 return (line[0].lower() + line[1:]).strip() | 222 return (line[0].lower() + line[1:]).strip() |
| 217 | 223 |
| 218 def execute(self, parser, args): | 224 def execute(self, parser, args): |
| 219 """Dispatches execution to the right command. | 225 """Dispatches execution to the right command. |
| 220 | 226 |
| (...skipping 24 matching lines...) Expand all Loading... |
| 245 return command(parser, args[1:]) | 251 return command(parser, args[1:]) |
| 246 | 252 |
| 247 cmdhelp = self.enumerate_commands().get('help') | 253 cmdhelp = self.enumerate_commands().get('help') |
| 248 if cmdhelp: | 254 if cmdhelp: |
| 249 # Not a known command. Default to help. | 255 # Not a known command. Default to help. |
| 250 self._add_command_usage(parser, cmdhelp) | 256 self._add_command_usage(parser, cmdhelp) |
| 251 return cmdhelp(parser, args) | 257 return cmdhelp(parser, args) |
| 252 | 258 |
| 253 # Nothing can be done. | 259 # Nothing can be done. |
| 254 return 2 | 260 return 2 |
| OLD | NEW |