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