1# Copyright 2018 The Abseil Authors. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15"""This module provides argparse integration with absl.flags. 16 17argparse_flags.ArgumentParser is a drop-in replacement for 18argparse.ArgumentParser. It takes care of collecting and defining absl flags 19in argparse. 20 21 22Here is a simple example: 23 24 # Assume the following absl.flags is defined in another module: 25 # 26 # from absl import flags 27 # flags.DEFINE_string('echo', None, 'The echo message.') 28 # 29 parser = argparse_flags.ArgumentParser( 30 description='A demo of absl.flags and argparse integration.') 31 parser.add_argument('--header', help='Header message to print.') 32 33 # The parser will also accept the absl flag `--echo`. 34 # The `header` value is available as `args.header` just like a regular 35 # argparse flag. The absl flag `--echo` continues to be available via 36 # `absl.flags.FLAGS` if you want to access it. 37 args = parser.parse_args() 38 39 # Example usages: 40 # ./program --echo='A message.' --header='A header' 41 # ./program --header 'A header' --echo 'A message.' 42 43 44Here is another example demonstrates subparsers: 45 46 parser = argparse_flags.ArgumentParser(description='A subcommands demo.') 47 parser.add_argument('--header', help='The header message to print.') 48 49 subparsers = parser.add_subparsers(help='The command to execute.') 50 51 roll_dice_parser = subparsers.add_parser( 52 'roll_dice', help='Roll a dice.', 53 # By default, absl flags can also be specified after the sub-command. 54 # To only allow them before sub-command, pass 55 # `inherited_absl_flags=None`. 56 inherited_absl_flags=None) 57 roll_dice_parser.add_argument('--num_faces', type=int, default=6) 58 roll_dice_parser.set_defaults(command=roll_dice) 59 60 shuffle_parser = subparsers.add_parser('shuffle', help='Shuffle inputs.') 61 shuffle_parser.add_argument( 62 'inputs', metavar='I', nargs='+', help='Inputs to shuffle.') 63 shuffle_parser.set_defaults(command=shuffle) 64 65 args = parser.parse_args(argv[1:]) 66 args.command(args) 67 68 # Example usages: 69 # ./program --echo='A message.' roll_dice --num_faces=6 70 # ./program shuffle --echo='A message.' 1 2 3 4 71 72 73There are several differences between absl.flags and argparse_flags: 74 751. Flags defined with absl.flags are parsed differently when using the 76 argparse parser. Notably: 77 78 1) absl.flags allows both single-dash and double-dash for any flag, and 79 doesn't distinguish them; argparse_flags only allows double-dash for 80 flag's regular name, and single-dash for flag's `short_name`. 81 2) Boolean flags in absl.flags can be specified with `--bool`, `--nobool`, 82 as well as `--bool=true/false` (though not recommended); 83 in argparse_flags, it only allows `--bool`, `--nobool`. 84 852. Help related flag differences: 86 1) absl.flags does not define help flags, absl.app does that; argparse_flags 87 defines help flags unless passed with `add_help=False`. 88 2) absl.app supports `--helpxml`; argparse_flags does not. 89 3) argparse_flags supports `-h`; absl.app does not. 90""" 91 92from __future__ import absolute_import 93from __future__ import division 94from __future__ import print_function 95 96import argparse 97import sys 98 99from absl import flags 100 101 102_BUILT_IN_FLAGS = frozenset({ 103 'help', 104 'helpshort', 105 'helpfull', 106 'helpxml', 107 'flagfile', 108 'undefok', 109}) 110 111 112class ArgumentParser(argparse.ArgumentParser): 113 """Custom ArgumentParser class to support special absl flags.""" 114 115 def __init__(self, **kwargs): 116 """Initializes ArgumentParser. 117 118 Args: 119 **kwargs: same as argparse.ArgumentParser, except: 120 1. It also accepts `inherited_absl_flags`: the absl flags to inherit. 121 The default is the global absl.flags.FLAGS instance. Pass None to 122 ignore absl flags. 123 2. The `prefix_chars` argument must be the default value '-'. 124 125 Raises: 126 ValueError: Raised when prefix_chars is not '-'. 127 """ 128 prefix_chars = kwargs.get('prefix_chars', '-') 129 if prefix_chars != '-': 130 raise ValueError( 131 'argparse_flags.ArgumentParser only supports "-" as the prefix ' 132 'character, found "{}".'.format(prefix_chars)) 133 134 # Remove inherited_absl_flags before calling super. 135 self._inherited_absl_flags = kwargs.pop('inherited_absl_flags', flags.FLAGS) 136 # Now call super to initialize argparse.ArgumentParser before calling 137 # add_argument in _define_absl_flags. 138 super(ArgumentParser, self).__init__(**kwargs) 139 140 if self.add_help: 141 # -h and --help are defined in super. 142 # Also add the --helpshort and --helpfull flags. 143 self.add_argument( 144 # Action 'help' defines a similar flag to -h/--help. 145 '--helpshort', action='help', 146 default=argparse.SUPPRESS, help=argparse.SUPPRESS) 147 self.add_argument( 148 '--helpfull', action=_HelpFullAction, 149 default=argparse.SUPPRESS, help='show full help message and exit') 150 151 if self._inherited_absl_flags: 152 self.add_argument('--undefok', help=argparse.SUPPRESS) 153 self._define_absl_flags(self._inherited_absl_flags) 154 155 def parse_known_args(self, args=None, namespace=None): 156 if args is None: 157 args = sys.argv[1:] 158 if self._inherited_absl_flags: 159 # Handle --flagfile. 160 # Explicitly specify force_gnu=True, since argparse behaves like 161 # gnu_getopt: flags can be specified after positional arguments. 162 args = self._inherited_absl_flags.read_flags_from_files( 163 args, force_gnu=True) 164 165 undefok_missing = object() 166 undefok = getattr(namespace, 'undefok', undefok_missing) 167 168 namespace, args = super(ArgumentParser, self).parse_known_args( 169 args, namespace) 170 171 # For Python <= 2.7.8: https://bugs.python.org/issue9351, a bug where 172 # sub-parsers don't preserve existing namespace attributes. 173 # Restore the undefok attribute if a sub-parser dropped it. 174 if undefok is not undefok_missing: 175 namespace.undefok = undefok 176 177 if self._inherited_absl_flags: 178 # Handle --undefok. At this point, `args` only contains unknown flags, 179 # so it won't strip defined flags that are also specified with --undefok. 180 # For Python <= 2.7.8: https://bugs.python.org/issue9351, a bug where 181 # sub-parsers don't preserve existing namespace attributes. The undefok 182 # attribute might not exist because a subparser dropped it. 183 if hasattr(namespace, 'undefok'): 184 args = _strip_undefok_args(namespace.undefok, args) 185 # absl flags are not exposed in the Namespace object. See Namespace: 186 # https://docs.python.org/3/library/argparse.html#argparse.Namespace. 187 del namespace.undefok 188 self._inherited_absl_flags.mark_as_parsed() 189 try: 190 self._inherited_absl_flags._assert_all_validators() # pylint: disable=protected-access 191 except flags.IllegalFlagValueError as e: 192 self.error(str(e)) 193 194 return namespace, args 195 196 def _define_absl_flags(self, absl_flags): 197 """Defines flags from absl_flags.""" 198 key_flags = set(absl_flags.get_key_flags_for_module(sys.argv[0])) 199 for name in absl_flags: 200 if name in _BUILT_IN_FLAGS: 201 # Do not inherit built-in flags. 202 continue 203 flag_instance = absl_flags[name] 204 # Each flags with short_name appears in FLAGS twice, so only define 205 # when the dictionary key is equal to the regular name. 206 if name == flag_instance.name: 207 # Suppress the flag in the help short message if it's not a main 208 # module's key flag. 209 suppress = flag_instance not in key_flags 210 self._define_absl_flag(flag_instance, suppress) 211 212 def _define_absl_flag(self, flag_instance, suppress): 213 """Defines a flag from the flag_instance.""" 214 flag_name = flag_instance.name 215 short_name = flag_instance.short_name 216 argument_names = ['--' + flag_name] 217 if short_name: 218 argument_names.insert(0, '-' + short_name) 219 if suppress: 220 helptext = argparse.SUPPRESS 221 else: 222 # argparse help string uses %-formatting. Escape the literal %'s. 223 helptext = flag_instance.help.replace('%', '%%') 224 if flag_instance.boolean: 225 # Only add the `no` form to the long name. 226 argument_names.append('--no' + flag_name) 227 self.add_argument( 228 *argument_names, action=_BooleanFlagAction, help=helptext, 229 metavar=flag_instance.name.upper(), 230 flag_instance=flag_instance) 231 else: 232 self.add_argument( 233 *argument_names, action=_FlagAction, help=helptext, 234 metavar=flag_instance.name.upper(), 235 flag_instance=flag_instance) 236 237 238class _FlagAction(argparse.Action): 239 """Action class for Abseil non-boolean flags.""" 240 241 def __init__(self, option_strings, dest, help, metavar, flag_instance): # pylint: disable=redefined-builtin 242 """Initializes _FlagAction. 243 244 Args: 245 option_strings: See argparse.Action. 246 dest: Ignored. The flag is always defined with dest=argparse.SUPPRESS. 247 help: See argparse.Action. 248 metavar: See argparse.Action. 249 flag_instance: absl.flags.Flag, the absl flag instance. 250 """ 251 del dest 252 self._flag_instance = flag_instance 253 super(_FlagAction, self).__init__( 254 option_strings=option_strings, 255 dest=argparse.SUPPRESS, 256 help=help, 257 metavar=metavar) 258 259 def __call__(self, parser, namespace, values, option_string=None): 260 """See https://docs.python.org/3/library/argparse.html#action-classes.""" 261 self._flag_instance.parse(values) 262 self._flag_instance.using_default_value = False 263 264 265class _BooleanFlagAction(argparse.Action): 266 """Action class for Abseil boolean flags.""" 267 268 def __init__(self, option_strings, dest, help, metavar, flag_instance): # pylint: disable=redefined-builtin 269 """Initializes _BooleanFlagAction. 270 271 Args: 272 option_strings: See argparse.Action. 273 dest: Ignored. The flag is always defined with dest=argparse.SUPPRESS. 274 help: See argparse.Action. 275 metavar: See argparse.Action. 276 flag_instance: absl.flags.Flag, the absl flag instance. 277 """ 278 del dest 279 self._flag_instance = flag_instance 280 flag_names = [self._flag_instance.name] 281 if self._flag_instance.short_name: 282 flag_names.append(self._flag_instance.short_name) 283 self._flag_names = frozenset(flag_names) 284 super(_BooleanFlagAction, self).__init__( 285 option_strings=option_strings, 286 dest=argparse.SUPPRESS, 287 nargs=0, # Does not accept values, only `--bool` or `--nobool`. 288 help=help, 289 metavar=metavar) 290 291 def __call__(self, parser, namespace, values, option_string=None): 292 """See https://docs.python.org/3/library/argparse.html#action-classes.""" 293 if not isinstance(values, list) or values: 294 raise ValueError('values must be an empty list.') 295 if option_string.startswith('--'): 296 option = option_string[2:] 297 else: 298 option = option_string[1:] 299 if option in self._flag_names: 300 self._flag_instance.parse('true') 301 else: 302 if not option.startswith('no') or option[2:] not in self._flag_names: 303 raise ValueError('invalid option_string: ' + option_string) 304 self._flag_instance.parse('false') 305 self._flag_instance.using_default_value = False 306 307 308class _HelpFullAction(argparse.Action): 309 """Action class for --helpfull flag.""" 310 311 def __init__(self, option_strings, dest, default, help): # pylint: disable=redefined-builtin 312 """Initializes _HelpFullAction. 313 314 Args: 315 option_strings: See argparse.Action. 316 dest: Ignored. The flag is always defined with dest=argparse.SUPPRESS. 317 default: Ignored. 318 help: See argparse.Action. 319 """ 320 del dest, default 321 super(_HelpFullAction, self).__init__( 322 option_strings=option_strings, 323 dest=argparse.SUPPRESS, 324 default=argparse.SUPPRESS, 325 nargs=0, 326 help=help) 327 328 def __call__(self, parser, namespace, values, option_string=None): 329 """See https://docs.python.org/3/library/argparse.html#action-classes.""" 330 # This only prints flags when help is not argparse.SUPPRESS. 331 # It includes user defined argparse flags, as well as main module's 332 # key absl flags. Other absl flags use argparse.SUPPRESS, so they aren't 333 # printed here. 334 parser.print_help() 335 336 absl_flags = parser._inherited_absl_flags # pylint: disable=protected-access 337 if absl_flags: 338 modules = sorted(absl_flags.flags_by_module_dict()) 339 main_module = sys.argv[0] 340 if main_module in modules: 341 # The main module flags are already printed in parser.print_help(). 342 modules.remove(main_module) 343 print(absl_flags._get_help_for_modules( # pylint: disable=protected-access 344 modules, prefix='', include_special_flags=True)) 345 parser.exit() 346 347 348def _strip_undefok_args(undefok, args): 349 """Returns a new list of args after removing flags in --undefok.""" 350 if undefok: 351 undefok_names = set(name.strip() for name in undefok.split(',')) 352 undefok_names |= set('no' + name for name in undefok_names) 353 # Remove undefok flags. 354 args = [arg for arg in args if not _is_undefok(arg, undefok_names)] 355 return args 356 357 358def _is_undefok(arg, undefok_names): 359 """Returns whether we can ignore arg based on a set of undefok flag names.""" 360 if not arg.startswith('-'): 361 return False 362 if arg.startswith('--'): 363 arg_without_dash = arg[2:] 364 else: 365 arg_without_dash = arg[1:] 366 if '=' in arg_without_dash: 367 name, _ = arg_without_dash.split('=', 1) 368 else: 369 name = arg_without_dash 370 if name in undefok_names: 371 return True 372 return False 373