Ben Murdoch | 097c5b2 | 2016-05-18 11:27:45 +0100 | [diff] [blame^] | 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 | """An option parser which handles the first arg as a command. |
| 6 | |
| 7 | Add other nice functionality such as printing a list of commands |
| 8 | and an example in usage. |
| 9 | """ |
| 10 | |
| 11 | import optparse |
| 12 | import sys |
| 13 | |
| 14 | |
| 15 | class CommandOptionParser(optparse.OptionParser): |
| 16 | """Wrapper class for OptionParser to help with listing commands.""" |
| 17 | |
| 18 | def __init__(self, *args, **kwargs): |
| 19 | """Creates a CommandOptionParser. |
| 20 | |
| 21 | Args: |
| 22 | commands_dict: A dictionary mapping command strings to an object defining |
| 23 | - add_options_func: Adds options to the option parser |
| 24 | - run_command_func: Runs the command itself. |
| 25 | example: An example command. |
| 26 | everything else: Passed to optparse.OptionParser contructor. |
| 27 | """ |
| 28 | self.commands_dict = kwargs.pop('commands_dict', {}) |
| 29 | self.example = kwargs.pop('example', '') |
| 30 | if not 'usage' in kwargs: |
| 31 | kwargs['usage'] = 'Usage: %prog <command> [options]' |
| 32 | optparse.OptionParser.__init__(self, *args, **kwargs) |
| 33 | |
| 34 | #override |
| 35 | def get_usage(self): |
| 36 | normal_usage = optparse.OptionParser.get_usage(self) |
| 37 | command_list = self.get_command_list() |
| 38 | example = self.get_example() |
| 39 | return self.expand_prog_name(normal_usage + example + command_list) |
| 40 | |
| 41 | #override |
| 42 | def get_command_list(self): |
| 43 | if self.commands_dict.keys(): |
| 44 | return '\nCommands:\n %s\n' % '\n '.join( |
| 45 | sorted(self.commands_dict.keys())) |
| 46 | return '' |
| 47 | |
| 48 | def get_example(self): |
| 49 | if self.example: |
| 50 | return '\nExample:\n %s\n' % self.example |
| 51 | return '' |
| 52 | |
| 53 | |
| 54 | def ParseAndExecute(option_parser, argv=None): |
| 55 | """Parses options/args from argv and runs the specified command. |
| 56 | |
| 57 | Args: |
| 58 | option_parser: A CommandOptionParser object. |
| 59 | argv: Command line arguments. If None, automatically draw from sys.argv. |
| 60 | |
| 61 | Returns: |
| 62 | An exit code. |
| 63 | """ |
| 64 | if not argv: |
| 65 | argv = sys.argv |
| 66 | |
| 67 | if len(argv) < 2 or argv[1] not in option_parser.commands_dict: |
| 68 | # Parse args first, if this is '--help', optparse will print help and exit |
| 69 | option_parser.parse_args(argv) |
| 70 | option_parser.error('Invalid command.') |
| 71 | |
| 72 | cmd = option_parser.commands_dict[argv[1]] |
| 73 | cmd.add_options_func(option_parser) |
| 74 | options, args = option_parser.parse_args(argv) |
| 75 | return cmd.run_command_func(argv[1], options, args, option_parser) |