from __future__ import print_function import abc import argparse from copy import deepcopy from functools import partial import six from pathlib2 import Path from trains_agent import definitions from trains_agent.session import Session HEADER = 'TRAINS-AGENT Deep Learning DevOps' class Parser(argparse.ArgumentParser): __default_subparser = None def __init__(self, usage_on_error=True, *args, **kwargs): super(Parser, self).__init__(fromfile_prefix_chars=definitions.FROM_FILE_PREFIX_CHARS, *args, **kwargs) self._usage_on_error = usage_on_error @property def choices(self): try: subparser = next( action for action in self._actions if isinstance(action, argparse._SubParsersAction)) except StopIteration: return {} return subparser.choices def error(self, message): if self._usage_on_error and message == argparse._('too few arguments'): self.print_help() print() self.exit(2, argparse._('%s: error: %s\n') % (self.prog, message)) super(Parser, self).error(message) def __getitem__(self, name): return self.choices[name] def remove_top_level_results(self, parse_results): """ Remove useless, artifact values :param parse_results: resulting namespace of parse_args, converted to dict ( vars(args) ) """ for action in self._actions: if action.dest != 'version': parse_results.pop(action.dest, None) for key in ('func', 'command', 'subcommand', 'action'): parse_results.pop(key, None) def set_default_subparser(self, name): self.__default_subparser = name def get_default_subparser(self): return self.choices[self.__default_subparser] def _parse_known_args(self, arg_strings, *args, **kwargs): in_args = set(arg_strings) d_sp = self.__default_subparser if d_sp is not None and not {'-h', '--help'}.intersection(in_args): for x in self._subparsers._actions: subparser_found = ( isinstance(x, argparse._SubParsersAction) and in_args.intersection(x._name_parser_map.keys()) ) if subparser_found: break else: # insert default in first position, this implies no # global options without a sub_parsers specified arg_strings = [d_sp] + arg_strings return super(Parser, self)._parse_known_args( arg_strings, *args, **kwargs ) class AliasedPseudoAction(argparse.Action): """ Action for choosing between sub-commands, including aliases """ def __init__(self, name, aliases, help): dest = name aliases = [a for a in aliases if a != name] if aliases: dest += ' (%s)' % ','.join(aliases) super(AliasedPseudoAction, self).__init__(option_strings=[], dest=dest, help=help) class AliasedSubParsersAction(argparse._SubParsersAction): """ Action for adding aliases for sub-commands """ def add_parser(self, name, **kwargs): aliases = kwargs.pop('aliases', []) parser = super(AliasedSubParsersAction, self).add_parser(name, **kwargs) # Make the aliases work for alias in aliases: self._name_parser_map[alias] = parser # Make the help text reflect them, first removing old help entry. help = kwargs.pop('help', None) if help: self._choices_actions.pop() pseudo_action = AliasedPseudoAction(name, aliases, help) self._choices_actions.append(pseudo_action) return parser class OnlyPluralChoicesHelpFormatter(argparse.HelpFormatter): @staticmethod def _metavar_formatter(action, default_metavar): if action.metavar is not None: result = action.metavar elif action.choices is not None: choice_strs = [str(choice) for choice in action.choices] choice_strs = [choice for choice in choice_strs if choice + 's' not in choice_strs] result = '{%s}' % ','.join(choice_strs) else: result = default_metavar def format(tuple_size): if isinstance(result, tuple): return result else: return (result, ) * tuple_size return format def hyphenate(s): return s.replace('_', '-') def add_args(parser, args): """ Add arguments to parser from args mapping :param parser: parser to add arguments to :type parser: argparse.ArgumentParser :param args: mapping of name -> other arguments to ArgumentParser.add_argument :type args: dict """ for arg_name, arg_params in args.items(): aliases = arg_params.pop('aliases', tuple()) parser.add_argument(arg_name, *aliases, **arg_params) def add_mutually_exclusive_groups(parser, groups): """ Add mutually exclusive groups to parser from list :param parser: parser to add groups to :param groups: list of dictionaries, each containing: 1. 'args': parameter to add_args 2. arguments to ArgumentParser.add_mutually_exclusive_group """ for group in groups: args = group.pop('args', {}) group_parser = parser.add_mutually_exclusive_group(**group) add_args(group_parser, args) def add_service(subparsers, name, commands, command_name_dest='command', formatter_class=argparse.RawDescriptionHelpFormatter, **kwargs): """ Add service commands to parser from arguments dictionary :param subparsers: subparsers object of ArgumentParser :param name: name of service :param commands: mapping of names to dictionaries, each of them containing: 1. 'args' - mapping of name -> other arguments to ArgumentParser.add_argument 2. 'help' - command description 3. 'mutually_exclusive_groups' - see add_mutually_exclusive_groups :param command_name_dest: name of attribute in which to store selected sub-command :param formatter_class; help formatter class :param kwargs: any other arguments to add_parser method of subparser object :return: service subparser """ commands = deepcopy(commands) service_parser = subparsers.add_parser( name, # aliases=(name.strip('s'),), formatter_class=formatter_class, **kwargs ) service_parser.register('action', 'parsers', AliasedSubParsersAction) service_parser.set_defaults(**{command_name_dest: name}) service_subparsers = service_parser.add_subparsers( title='{} commands'.format(name.capitalize()), parser_class=partial(Parser, usage_on_error=False), dest='action') # This is a fix for a bug in python3's argparse: running "trains-agent some_service" fails service_subparsers.required = True for name, subparser in commands.pop('subparsers', {}).items(): add_service(service_subparsers, name, command_name_dest='subcommand', **subparser) for command_name, command in commands.items(): command_type = command.pop('type', None) mutually_exclusive_groups = command.pop('mutually_exclusive_groups', []) func = command.pop('func', command_name) args = command.pop('args', {}) command_parser = service_subparsers.add_parser(hyphenate(command_name), **command) if command_type: command_type.make(command_parser) command_parser.set_defaults(func=func) add_mutually_exclusive_groups(command_parser, mutually_exclusive_groups) add_args(command_parser, args) return service_parser @six.add_metaclass(abc.ABCMeta) class CommandType(object): def __init__(self, *args, **kwargs): self._args = args self._kwargs = kwargs def make(self, parser): return self._make(parser, *self._args, **self._kwargs) @abc.abstractmethod def _make(self, *args, **kwargs): pass class ListCommand(CommandType): @staticmethod def _make(parser, default_value, pagination=False, tree=False): if tree: table_group = parser.add_mutually_exclusive_group() tree_action = table_group.add_argument( '--tree', help='Tree view output', action='store_true', default=False) csv_group = parser.add_mutually_exclusive_group() # hack: tree cannot be used with either csv or table, which can be used # which each other csv_group._group_actions.append(tree_action) else: csv_group = table_group = parser table_group.add_argument( '--table', help='Select table columns ("#" separated, default: %(default)s)', default=default_value) csv_group.add_argument( '--csv', help='Generate CSV output to specified path', default=None) parser.add_argument( '--no-headers', action='store_false', dest='headers', help='Do not print table/csv headers') parser.add_argument( '--sort', help='Fields to sort by (same format as --table)') parser.add_argument( '--ascending', default=None, dest='sort_reverse', action='store_false', help='Sort in ascending order (default)') parser.add_argument( '--descending', default=None, dest='sort_reverse', action='store_true', help='Sort in descending order') if not pagination: return parser.add_argument( '--page', help='Page number to show (default: %(default)s)', default=0, type=bound_number_type(minimum=0)) parser.add_argument( '--page-size', help='Size of page (default: %(default)s)', default=50, type=bound_number_type(minimum=1)) parser.add_argument( '--no-pagination', action='store_false', dest='pagination', help='Disable pagination (return all results)') class _HelpAction(argparse._HelpAction): def __call__(self, parser, namespace, values, option_string=None): # print header print(HEADER + '\n') parser.print_help() print('') parser.exit() class _DetailedHelpAction(argparse._HelpAction): def __call__(self, parser, namespace, values, option_string=None): # print header print(HEADER + '\n') parser.print_help() print('\n') # retrieve subparsers from parser subparsers_actions = [ action for action in parser._actions if isinstance(action, argparse._SubParsersAction) ] # iterate and print help for each suparser for subparsers_action in subparsers_actions: # get all subparsers and print help for choice, subparsercmd in subparsers_action.choices.items(): # split help into lines so we can skip the header text = subparsercmd.format_help().split('\n') # find first line of command (skip usage and header) for i, t in enumerate(text): if t.startswith(choice.title()): break # print help command prefix print(text[i]) # print help per sub-commands, we actually assume only one subact = [ action for action in subparsercmd._actions if isinstance(action, argparse._SubParsersAction) ] # per action print all parameters for j, t in enumerate(text[i + 2:]): print(t) k = t.split() if not k: continue try: subc = subact[0].choices[k[0]] except KeyError: continue # hack so we can control formatting in one place # otherwise we need to update all # the parsers when we create them subc.formatter_class = lambda prog: argparse.HelpFormatter(prog, width=120) subchelp = subc.format_help().split('\n') # skip until we reach "optional arguments:" for si, st in enumerate(subchelp): if st.startswith('optional arguments:'): break # print help with single tab indent for st in subchelp[si + 1:]: print('\t %s' % st) parser.exit() def base_arguments(top_parser): top_parser.register('action', 'parsers', AliasedSubParsersAction) top_parser.add_argument('-h', action=_HelpAction, help='Displays summary of all commands') top_parser.add_argument( '--help', action=_DetailedHelpAction, help='Detailed help of command line interface') top_parser.add_argument( '--version', action='version', version='TRAINS-AGENT version %s' % Session.version, help='TRAINS-AGENT version number') top_parser.add_argument( '--config-file', help='Use a different configuration file (default: "{}")'.format(definitions.CONFIG_FILE)) top_parser.add_argument('--debug', '-d', action='store_true', help='print debug information') top_parser.add_argument( '--trace', '-t', action='store_true', help='Trace execution to a temporary file.', ) def bound_number_type(minimum=None, maximum=None): """ bound_number_type Creates a bounded integer "type" (validator function) for use with argparse.ArgumentParser.add_argument. At least one of ``minimum`` and ``maximum`` must be passed. :param minimum: maximum allowed value :param maximum: minimum allowed value """ if minimum is maximum is None: raise ValueError('either "minimum" or "maximum" must be provided') def bound_int(arg): num = int(arg) if minimum is not None and num < minimum: raise argparse.ArgumentTypeError('minimum value is {}'.format(minimum)) if maximum is not None and num > maximum: raise argparse.ArgumentTypeError('maximum value is {}'.format(minimum)) return num return bound_int def real_path_type(string): path = Path(string).expanduser() if not path.exists(): raise argparse.ArgumentTypeError('"{}": No such file or directory'.format(path)) return path class ObjectID(object): def __init__(self, name, service=None): self.name = name self.service = service def foreign_object_id(service): return partial(ObjectID, service=service)