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