JIRA: BOTTLENECKS-29
[bottlenecks.git] / vstf / vstf / common / vstfcli.py
1 ##############################################################################
2 # Copyright (c) 2015 Huawei Technologies Co.,Ltd and others.
3 #
4 # All rights reserved. This program and the accompanying materials
5 # are made available under the terms of the Apache License, Version 2.0
6 # which accompanies this distribution, and is available at
7 # http://www.apache.org/licenses/LICENSE-2.0
8 ##############################################################################
9
10 import argparse
11 import sys
12
13
14 class VstfHelpFormatter(argparse.HelpFormatter):
15     def start_section(self, heading):
16         # Title-case the headings
17         heading = '%s%s' % (heading[0].upper(), heading[1:])
18         super(VstfHelpFormatter, self).start_section(heading)
19
20
21 class VstfParser(argparse.ArgumentParser):
22     def __init__(self,
23                  prog='vstf',
24                  description="",
25                  epilog='',
26                  add_help=True,
27                  formatter_class=VstfHelpFormatter):
28
29         super(VstfParser, self).__init__(
30             prog=prog,
31             description=description,
32             epilog=epilog,
33             add_help=add_help,
34             formatter_class=formatter_class)
35         self.subcommands = {}
36
37     def _find_actions(self, subparsers, actions_module):
38         for attr in (a for a in dir(actions_module) if a.startswith('do_')):
39             command = attr[3:].replace('_', '-')
40             callback = getattr(actions_module, attr)
41             desc = callback.__doc__ or ''
42             action_help = desc.strip()
43             arguments = getattr(callback, 'arguments', [])
44             subparser = subparsers.add_parser(command,
45                                               help=action_help,
46                                               description=desc,
47                                               add_help=False,
48                                               formatter_class=VstfHelpFormatter)
49             subparser.add_argument('-h', '--help',
50                                    action='help',
51                                    help=argparse.SUPPRESS)
52             self.subcommands[command] = subparser
53             for (args, kwargs) in arguments:
54                 subparser.add_argument(*args, **kwargs)
55             subparser.set_defaults(func=callback)
56
57     def set_subcommand_parser(self, target, metavar="<subcommand>"):
58         subparsers = self.add_subparsers(metavar=metavar)
59         self._find_actions(subparsers, target)
60         return subparsers
61
62     def set_parser_to_subcommand(self, subparser, target):
63         self._find_actions(subparser, target)
64
65
66 if __name__ == "__main__":
67     from vstf.common import test_func
68     parser = VstfParser(prog="vstf", description="test parser")
69     parser.set_subcommand_parser(test_func)
70     args = parser.parse_args(sys.argv[1:])
71     args.func(args)