1# This Source Code Form is subject to the terms of the Mozilla Public 2# License, v. 2.0. If a copy of the MPL was not distributed with this 3# file, You can obtain one at http://mozilla.org/MPL/2.0/. 4 5from __future__ import absolute_import, unicode_literals 6 7import argparse 8import difflib 9import shlex 10import sys 11 12from operator import itemgetter 13 14from .base import ( 15 NoCommandError, 16 UnknownCommandError, 17 UnrecognizedArgumentError, 18) 19from .decorators import SettingsProvider 20 21 22@SettingsProvider 23class DispatchSettings(): 24 config_settings = [ 25 ('alias.*', 'string', """ 26Create a command alias of the form `<alias>=<command> <args>`. 27Aliases can also be used to set default arguments: 28<command>=<command> <args> 29""".strip()), 30 ] 31 32 33class CommandFormatter(argparse.HelpFormatter): 34 """Custom formatter to format just a subcommand.""" 35 36 def add_usage(self, *args): 37 pass 38 39 40class CommandAction(argparse.Action): 41 """An argparse action that handles mach commands. 42 43 This class is essentially a reimplementation of argparse's sub-parsers 44 feature. We first tried to use sub-parsers. However, they were missing 45 features like grouping of commands (http://bugs.python.org/issue14037). 46 47 The way this works involves light magic and a partial understanding of how 48 argparse works. 49 50 Arguments registered with an argparse.ArgumentParser have an action 51 associated with them. An action is essentially a class that when called 52 does something with the encountered argument(s). This class is one of those 53 action classes. 54 55 An instance of this class is created doing something like: 56 57 parser.add_argument('command', action=CommandAction, registrar=r) 58 59 Note that a mach.registrar.Registrar instance is passed in. The Registrar 60 holds information on all the mach commands that have been registered. 61 62 When this argument is registered with the ArgumentParser, an instance of 63 this class is instantiated. One of the subtle but important things it does 64 is tell the argument parser that it's interested in *all* of the remaining 65 program arguments. So, when the ArgumentParser calls this action, we will 66 receive the command name plus all of its arguments. 67 68 For more, read the docs in __call__. 69 """ 70 def __init__(self, option_strings, dest, required=True, default=None, 71 registrar=None, context=None): 72 # A proper API would have **kwargs here. However, since we are a little 73 # hacky, we intentionally omit it as a way of detecting potentially 74 # breaking changes with argparse's implementation. 75 # 76 # In a similar vein, default is passed in but is not needed, so we drop 77 # it. 78 argparse.Action.__init__(self, option_strings, dest, required=required, 79 help=argparse.SUPPRESS, nargs=argparse.REMAINDER) 80 81 self._mach_registrar = registrar 82 self._context = context 83 84 def __call__(self, parser, namespace, values, option_string=None): 85 """This is called when the ArgumentParser has reached our arguments. 86 87 Since we always register ourselves with nargs=argparse.REMAINDER, 88 values should be a list of remaining arguments to parse. The first 89 argument should be the name of the command to invoke and all remaining 90 arguments are arguments for that command. 91 92 The gist of the flow is that we look at the command being invoked. If 93 it's *help*, we handle that specially (because argparse's default help 94 handler isn't satisfactory). Else, we create a new, independent 95 ArgumentParser instance for just the invoked command (based on the 96 information contained in the command registrar) and feed the arguments 97 into that parser. We then merge the results with the main 98 ArgumentParser. 99 """ 100 if namespace.help: 101 # -h or --help is in the global arguments. 102 self._handle_main_help(parser, namespace.verbose) 103 sys.exit(0) 104 elif values: 105 command = values[0].lower() 106 args = values[1:] 107 if command == 'help': 108 if args and args[0] not in ['-h', '--help']: 109 # Make sure args[0] is indeed a command. 110 self._handle_command_help(parser, args[0], args) 111 else: 112 self._handle_main_help(parser, namespace.verbose) 113 sys.exit(0) 114 elif '-h' in args or '--help' in args: 115 # -h or --help is in the command arguments. 116 if '--' in args: 117 # -- is in command arguments 118 if '-h' in args[:args.index('--')] or '--help' in args[:args.index('--')]: 119 # Honor -h or --help only if it appears before -- 120 self._handle_command_help(parser, command, args) 121 sys.exit(0) 122 else: 123 self._handle_command_help(parser, command, args) 124 sys.exit(0) 125 else: 126 raise NoCommandError() 127 128 # First see if the this is a user-defined alias 129 if command in self._context.settings.alias: 130 alias = self._context.settings.alias[command] 131 defaults = shlex.split(alias) 132 command = defaults.pop(0) 133 args = defaults + args 134 135 if command not in self._mach_registrar.command_handlers: 136 # Try to find similar commands, may raise UnknownCommandError. 137 command = self._suggest_command(command) 138 139 handler = self._mach_registrar.command_handlers.get(command) 140 141 usage = '%(prog)s [global arguments] ' + command + \ 142 ' [command arguments]' 143 144 subcommand = None 145 146 # If there are sub-commands, parse the intent out immediately. 147 if handler.subcommand_handlers and args: 148 # mach <command> help <subcommand> 149 if set(args).intersection(('help', '--help')): 150 self._handle_subcommand_help(parser, handler, args) 151 sys.exit(0) 152 # mach <command> <subcommand> ... 153 elif args[0] in handler.subcommand_handlers: 154 subcommand = args[0] 155 handler = handler.subcommand_handlers[subcommand] 156 usage = '%(prog)s [global arguments] ' + command + ' ' + \ 157 subcommand + ' [command arguments]' 158 args.pop(0) 159 160 # We create a new parser, populate it with the command's arguments, 161 # then feed all remaining arguments to it, merging the results 162 # with ourselves. This is essentially what argparse subparsers 163 # do. 164 165 parser_args = { 166 'add_help': False, 167 'usage': usage, 168 } 169 170 remainder = None 171 172 if handler.parser: 173 subparser = handler.parser 174 subparser.context = self._context 175 for arg in subparser._actions[:]: 176 if arg.nargs == argparse.REMAINDER: 177 subparser._actions.remove(arg) 178 remainder = (arg.dest,), {'default': arg.default, 179 'nargs': arg.nargs, 180 'help': arg.help} 181 else: 182 subparser = argparse.ArgumentParser(**parser_args) 183 184 for arg in handler.arguments: 185 # Remove our group keyword; it's not needed here. 186 group_name = arg[1].get('group') 187 if group_name: 188 del arg[1]['group'] 189 190 if arg[1].get('nargs') == argparse.REMAINDER: 191 # parse_known_args expects all argparse.REMAINDER ('...') 192 # arguments to be all stuck together. Instead, we want them to 193 # pick any extra argument, wherever they are. 194 # Assume a limited CommandArgument for those arguments. 195 assert len(arg[0]) == 1 196 assert all(k in ('default', 'nargs', 'help') for k in arg[1]) 197 remainder = arg 198 else: 199 subparser.add_argument(*arg[0], **arg[1]) 200 201 # We define the command information on the main parser result so as to 202 # not interfere with arguments passed to the command. 203 setattr(namespace, 'mach_handler', handler) 204 setattr(namespace, 'command', command) 205 setattr(namespace, 'subcommand', subcommand) 206 207 command_namespace, extra = subparser.parse_known_args(args) 208 setattr(namespace, 'command_args', command_namespace) 209 if remainder: 210 (name,), options = remainder 211 # parse_known_args usefully puts all arguments after '--' in 212 # extra, but also puts '--' there. We don't want to pass it down 213 # to the command handler. Note that if multiple '--' are on the 214 # command line, only the first one is removed, so that subsequent 215 # ones are passed down. 216 if '--' in extra: 217 extra.remove('--') 218 219 # Commands with argparse.REMAINDER arguments used to force the 220 # other arguments to be '+' prefixed. If a user now passes such 221 # an argument, if will silently end up in extra. So, check if any 222 # of the allowed arguments appear in a '+' prefixed form, and error 223 # out if that's the case. 224 for args, _ in handler.arguments: 225 for arg in args: 226 arg = arg.replace('-', '+', 1) 227 if arg in extra: 228 raise UnrecognizedArgumentError(command, [arg]) 229 230 if extra: 231 setattr(command_namespace, name, extra) 232 else: 233 setattr(command_namespace, name, options.get('default', [])) 234 elif extra and handler.cls.__name__ != 'DeprecatedCommands': 235 raise UnrecognizedArgumentError(command, extra) 236 237 def _handle_main_help(self, parser, verbose): 238 # Since we don't need full sub-parser support for the main help output, 239 # we create groups in the ArgumentParser and populate each group with 240 # arguments corresponding to command names. This has the side-effect 241 # that argparse renders it nicely. 242 r = self._mach_registrar 243 disabled_commands = [] 244 245 cats = [(k, v[2]) for k, v in r.categories.items()] 246 sorted_cats = sorted(cats, key=itemgetter(1), reverse=True) 247 for category, priority in sorted_cats: 248 group = None 249 250 for command in sorted(r.commands_by_category[category]): 251 handler = r.command_handlers[command] 252 253 # Instantiate a handler class to see if it should be filtered 254 # out for the current context or not. Condition functions can be 255 # applied to the command's decorator. 256 if handler.conditions: 257 if handler.pass_context: 258 instance = handler.cls(self._context) 259 else: 260 instance = handler.cls() 261 262 is_filtered = False 263 for c in handler.conditions: 264 if not c(instance): 265 is_filtered = True 266 break 267 if is_filtered: 268 description = handler.description 269 disabled_command = {'command': command, 'description': description} 270 disabled_commands.append(disabled_command) 271 continue 272 273 if group is None: 274 title, description, _priority = r.categories[category] 275 group = parser.add_argument_group(title, description) 276 277 description = handler.description 278 group.add_argument(command, help=description, 279 action='store_true') 280 281 if disabled_commands and 'disabled' in r.categories: 282 title, description, _priority = r.categories['disabled'] 283 group = parser.add_argument_group(title, description) 284 if verbose: 285 for c in disabled_commands: 286 group.add_argument(c['command'], help=c['description'], 287 action='store_true') 288 289 parser.print_help() 290 291 def _populate_command_group(self, parser, handler, group): 292 extra_groups = {} 293 for group_name in handler.argument_group_names: 294 group_full_name = 'Command Arguments for ' + group_name 295 extra_groups[group_name] = \ 296 parser.add_argument_group(group_full_name) 297 298 for arg in handler.arguments: 299 # Apply our group keyword. 300 group_name = arg[1].get('group') 301 if group_name: 302 del arg[1]['group'] 303 group = extra_groups[group_name] 304 group.add_argument(*arg[0], **arg[1]) 305 306 def _handle_command_help(self, parser, command, args): 307 handler = self._mach_registrar.command_handlers.get(command) 308 309 if not handler: 310 raise UnknownCommandError(command, 'query') 311 312 if handler.subcommand_handlers: 313 self._handle_subcommand_help(parser, handler, args) 314 return 315 316 # This code is worth explaining. Because we are doing funky things with 317 # argument registration to allow the same option in both global and 318 # command arguments, we can't simply put all arguments on the same 319 # parser instance because argparse would complain. We can't register an 320 # argparse subparser here because it won't properly show help for 321 # global arguments. So, we employ a strategy similar to command 322 # execution where we construct a 2nd, independent ArgumentParser for 323 # just the command data then supplement the main help's output with 324 # this 2nd parser's. We use a custom formatter class to ignore some of 325 # the help output. 326 parser_args = { 327 'formatter_class': CommandFormatter, 328 'add_help': False, 329 } 330 331 if handler.parser: 332 c_parser = handler.parser 333 c_parser.context = self._context 334 c_parser.formatter_class = NoUsageFormatter 335 # Accessing _action_groups is a bit shady. We are highly dependent 336 # on the argparse implementation not changing. We fail fast to 337 # detect upstream changes so we can intelligently react to them. 338 group = c_parser._action_groups[1] 339 340 # By default argparse adds two groups called "positional arguments" 341 # and "optional arguments". We want to rename these to reflect standard 342 # mach terminology. 343 c_parser._action_groups[0].title = 'Command Parameters' 344 c_parser._action_groups[1].title = 'Command Arguments' 345 346 if not handler.description: 347 handler.description = c_parser.description 348 c_parser.description = None 349 else: 350 c_parser = argparse.ArgumentParser(**parser_args) 351 group = c_parser.add_argument_group('Command Arguments') 352 353 self._populate_command_group(c_parser, handler, group) 354 355 # Set the long help of the command to the docstring (if present) or 356 # the command decorator description argument (if present). 357 if handler.docstring: 358 parser.description = format_docstring(handler.docstring) 359 elif handler.description: 360 parser.description = handler.description 361 362 parser.usage = '%(prog)s [global arguments] ' + command + \ 363 ' [command arguments]' 364 365 # This is needed to preserve line endings in the description field, 366 # which may be populated from a docstring. 367 parser.formatter_class = argparse.RawDescriptionHelpFormatter 368 parser.print_help() 369 print('') 370 c_parser.print_help() 371 372 def _handle_subcommand_main_help(self, parser, handler): 373 parser.usage = '%(prog)s [global arguments] ' + handler.name + \ 374 ' subcommand [subcommand arguments]' 375 group = parser.add_argument_group('Sub Commands') 376 377 for subcommand, subhandler in sorted(handler.subcommand_handlers.iteritems()): 378 group.add_argument(subcommand, help=subhandler.description, 379 action='store_true') 380 381 if handler.docstring: 382 parser.description = format_docstring(handler.docstring) 383 384 parser.formatter_class = argparse.RawDescriptionHelpFormatter 385 386 parser.print_help() 387 388 def _handle_subcommand_help(self, parser, handler, args): 389 subcommand = set(args).intersection(handler.subcommand_handlers.keys()) 390 if not subcommand: 391 return self._handle_subcommand_main_help(parser, handler) 392 393 subcommand = subcommand.pop() 394 subhandler = handler.subcommand_handlers[subcommand] 395 396 c_parser = subhandler.parser or argparse.ArgumentParser(add_help=False) 397 c_parser.formatter_class = CommandFormatter 398 399 group = c_parser.add_argument_group('Sub Command Arguments') 400 self._populate_command_group(c_parser, subhandler, group) 401 402 if subhandler.docstring: 403 parser.description = format_docstring(subhandler.docstring) 404 405 parser.formatter_class = argparse.RawDescriptionHelpFormatter 406 parser.usage = '%(prog)s [global arguments] ' + handler.name + \ 407 ' ' + subcommand + ' [command arguments]' 408 409 parser.print_help() 410 print('') 411 c_parser.print_help() 412 413 def _suggest_command(self, command): 414 # Make sure we don't suggest any deprecated commands. 415 names = [h.name for h in self._mach_registrar.command_handlers.values() 416 if h.cls.__name__ != 'DeprecatedCommands'] 417 # We first try to look for a valid command that is very similar to the given command. 418 suggested_commands = difflib.get_close_matches(command, names, cutoff=0.8) 419 # If we find more than one matching command, or no command at all, 420 # we give command suggestions instead (with a lower matching threshold). 421 # All commands that start with the given command (for instance: 422 # 'mochitest-plain', 'mochitest-chrome', etc. for 'mochitest-') 423 # are also included. 424 if len(suggested_commands) != 1: 425 suggested_commands = set(difflib.get_close_matches(command, names, cutoff=0.5)) 426 suggested_commands |= {cmd for cmd in names if cmd.startswith(command)} 427 raise UnknownCommandError(command, 'run', suggested_commands) 428 sys.stderr.write("We're assuming the '%s' command is '%s' and we're " 429 "executing it for you.\n\n" % (command, suggested_commands[0])) 430 return suggested_commands[0] 431 432 433class NoUsageFormatter(argparse.HelpFormatter): 434 def _format_usage(self, *args, **kwargs): 435 return "" 436 437 438def format_docstring(docstring): 439 """Format a raw docstring into something suitable for presentation. 440 441 This function is based on the example function in PEP-0257. 442 """ 443 if not docstring: 444 return '' 445 lines = docstring.expandtabs().splitlines() 446 indent = sys.maxint 447 for line in lines[1:]: 448 stripped = line.lstrip() 449 if stripped: 450 indent = min(indent, len(line) - len(stripped)) 451 trimmed = [lines[0].strip()] 452 if indent < sys.maxint: 453 for line in lines[1:]: 454 trimmed.append(line[indent:].rstrip()) 455 while trimmed and not trimmed[-1]: 456 trimmed.pop() 457 while trimmed and not trimmed[0]: 458 trimmed.pop(0) 459 return '\n'.join(trimmed) 460