Spaces:
Running
Running
# Copyright 2018 The Abseil Authors. | |
# | |
# Licensed under the Apache License, Version 2.0 (the "License"); | |
# you may not use this file except in compliance with the License. | |
# You may obtain a copy of the License at | |
# | |
# http://www.apache.org/licenses/LICENSE-2.0 | |
# | |
# Unless required by applicable law or agreed to in writing, software | |
# distributed under the License is distributed on an "AS IS" BASIS, | |
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
# See the License for the specific language governing permissions and | |
# limitations under the License. | |
"""This module provides argparse integration with absl.flags. | |
``argparse_flags.ArgumentParser`` is a drop-in replacement for | |
:class:`argparse.ArgumentParser`. It takes care of collecting and defining absl | |
flags in :mod:`argparse`. | |
Here is a simple example:: | |
# Assume the following absl.flags is defined in another module: | |
# | |
# from absl import flags | |
# flags.DEFINE_string('echo', None, 'The echo message.') | |
# | |
parser = argparse_flags.ArgumentParser( | |
description='A demo of absl.flags and argparse integration.') | |
parser.add_argument('--header', help='Header message to print.') | |
# The parser will also accept the absl flag `--echo`. | |
# The `header` value is available as `args.header` just like a regular | |
# argparse flag. The absl flag `--echo` continues to be available via | |
# `absl.flags.FLAGS` if you want to access it. | |
args = parser.parse_args() | |
# Example usages: | |
# ./program --echo='A message.' --header='A header' | |
# ./program --header 'A header' --echo 'A message.' | |
Here is another example demonstrates subparsers:: | |
parser = argparse_flags.ArgumentParser(description='A subcommands demo.') | |
parser.add_argument('--header', help='The header message to print.') | |
subparsers = parser.add_subparsers(help='The command to execute.') | |
roll_dice_parser = subparsers.add_parser( | |
'roll_dice', help='Roll a dice.', | |
# By default, absl flags can also be specified after the sub-command. | |
# To only allow them before sub-command, pass | |
# `inherited_absl_flags=None`. | |
inherited_absl_flags=None) | |
roll_dice_parser.add_argument('--num_faces', type=int, default=6) | |
roll_dice_parser.set_defaults(command=roll_dice) | |
shuffle_parser = subparsers.add_parser('shuffle', help='Shuffle inputs.') | |
shuffle_parser.add_argument( | |
'inputs', metavar='I', nargs='+', help='Inputs to shuffle.') | |
shuffle_parser.set_defaults(command=shuffle) | |
args = parser.parse_args(argv[1:]) | |
args.command(args) | |
# Example usages: | |
# ./program --echo='A message.' roll_dice --num_faces=6 | |
# ./program shuffle --echo='A message.' 1 2 3 4 | |
There are several differences between :mod:`absl.flags` and | |
:mod:`~absl.flags.argparse_flags`: | |
1. Flags defined with absl.flags are parsed differently when using the | |
argparse parser. Notably: | |
1) absl.flags allows both single-dash and double-dash for any flag, and | |
doesn't distinguish them; argparse_flags only allows double-dash for | |
flag's regular name, and single-dash for flag's ``short_name``. | |
2) Boolean flags in absl.flags can be specified with ``--bool``, | |
``--nobool``, as well as ``--bool=true/false`` (though not recommended); | |
in argparse_flags, it only allows ``--bool``, ``--nobool``. | |
2. Help related flag differences: | |
1) absl.flags does not define help flags, absl.app does that; argparse_flags | |
defines help flags unless passed with ``add_help=False``. | |
2) absl.app supports ``--helpxml``; argparse_flags does not. | |
3) argparse_flags supports ``-h``; absl.app does not. | |
""" | |
import argparse | |
import sys | |
from absl import flags | |
_BUILT_IN_FLAGS = frozenset({ | |
'help', | |
'helpshort', | |
'helpfull', | |
'helpxml', | |
'flagfile', | |
'undefok', | |
}) | |
class ArgumentParser(argparse.ArgumentParser): | |
"""Custom ArgumentParser class to support special absl flags.""" | |
def __init__(self, **kwargs): | |
"""Initializes ArgumentParser. | |
Args: | |
**kwargs: same as argparse.ArgumentParser, except: | |
1. It also accepts `inherited_absl_flags`: the absl flags to inherit. | |
The default is the global absl.flags.FLAGS instance. Pass None to | |
ignore absl flags. | |
2. The `prefix_chars` argument must be the default value '-'. | |
Raises: | |
ValueError: Raised when prefix_chars is not '-'. | |
""" | |
prefix_chars = kwargs.get('prefix_chars', '-') | |
if prefix_chars != '-': | |
raise ValueError( | |
'argparse_flags.ArgumentParser only supports "-" as the prefix ' | |
'character, found "{}".'.format(prefix_chars)) | |
# Remove inherited_absl_flags before calling super. | |
self._inherited_absl_flags = kwargs.pop('inherited_absl_flags', flags.FLAGS) | |
# Now call super to initialize argparse.ArgumentParser before calling | |
# add_argument in _define_absl_flags. | |
super(ArgumentParser, self).__init__(**kwargs) | |
if self.add_help: | |
# -h and --help are defined in super. | |
# Also add the --helpshort and --helpfull flags. | |
self.add_argument( | |
# Action 'help' defines a similar flag to -h/--help. | |
'--helpshort', action='help', | |
default=argparse.SUPPRESS, help=argparse.SUPPRESS) | |
self.add_argument( | |
'--helpfull', action=_HelpFullAction, | |
default=argparse.SUPPRESS, help='show full help message and exit') | |
if self._inherited_absl_flags is not None: | |
self.add_argument( | |
'--undefok', default=argparse.SUPPRESS, help=argparse.SUPPRESS) | |
self._define_absl_flags(self._inherited_absl_flags) | |
def parse_known_args(self, args=None, namespace=None): | |
if args is None: | |
args = sys.argv[1:] | |
if self._inherited_absl_flags is not None: | |
# Handle --flagfile. | |
# Explicitly specify force_gnu=True, since argparse behaves like | |
# gnu_getopt: flags can be specified after positional arguments. | |
args = self._inherited_absl_flags.read_flags_from_files( | |
args, force_gnu=True) | |
undefok_missing = object() | |
undefok = getattr(namespace, 'undefok', undefok_missing) | |
namespace, args = super(ArgumentParser, self).parse_known_args( | |
args, namespace) | |
# For Python <= 2.7.8: https://bugs.python.org/issue9351, a bug where | |
# sub-parsers don't preserve existing namespace attributes. | |
# Restore the undefok attribute if a sub-parser dropped it. | |
if undefok is not undefok_missing: | |
namespace.undefok = undefok | |
if self._inherited_absl_flags is not None: | |
# Handle --undefok. At this point, `args` only contains unknown flags, | |
# so it won't strip defined flags that are also specified with --undefok. | |
# For Python <= 2.7.8: https://bugs.python.org/issue9351, a bug where | |
# sub-parsers don't preserve existing namespace attributes. The undefok | |
# attribute might not exist because a subparser dropped it. | |
if hasattr(namespace, 'undefok'): | |
args = _strip_undefok_args(namespace.undefok, args) | |
# absl flags are not exposed in the Namespace object. See Namespace: | |
# https://docs.python.org/3/library/argparse.html#argparse.Namespace. | |
del namespace.undefok | |
self._inherited_absl_flags.mark_as_parsed() | |
try: | |
self._inherited_absl_flags.validate_all_flags() | |
except flags.IllegalFlagValueError as e: | |
self.error(str(e)) | |
return namespace, args | |
def _define_absl_flags(self, absl_flags): | |
"""Defines flags from absl_flags.""" | |
key_flags = set(absl_flags.get_key_flags_for_module(sys.argv[0])) | |
for name in absl_flags: | |
if name in _BUILT_IN_FLAGS: | |
# Do not inherit built-in flags. | |
continue | |
flag_instance = absl_flags[name] | |
# Each flags with short_name appears in FLAGS twice, so only define | |
# when the dictionary key is equal to the regular name. | |
if name == flag_instance.name: | |
# Suppress the flag in the help short message if it's not a main | |
# module's key flag. | |
suppress = flag_instance not in key_flags | |
self._define_absl_flag(flag_instance, suppress) | |
def _define_absl_flag(self, flag_instance, suppress): | |
"""Defines a flag from the flag_instance.""" | |
flag_name = flag_instance.name | |
short_name = flag_instance.short_name | |
argument_names = ['--' + flag_name] | |
if short_name: | |
argument_names.insert(0, '-' + short_name) | |
if suppress: | |
helptext = argparse.SUPPRESS | |
else: | |
# argparse help string uses %-formatting. Escape the literal %'s. | |
helptext = flag_instance.help.replace('%', '%%') | |
if flag_instance.boolean: | |
# Only add the `no` form to the long name. | |
argument_names.append('--no' + flag_name) | |
self.add_argument( | |
*argument_names, action=_BooleanFlagAction, help=helptext, | |
metavar=flag_instance.name.upper(), | |
flag_instance=flag_instance) | |
else: | |
self.add_argument( | |
*argument_names, action=_FlagAction, help=helptext, | |
metavar=flag_instance.name.upper(), | |
flag_instance=flag_instance) | |
class _FlagAction(argparse.Action): | |
"""Action class for Abseil non-boolean flags.""" | |
def __init__( | |
self, | |
option_strings, | |
dest, | |
help, # pylint: disable=redefined-builtin | |
metavar, | |
flag_instance, | |
default=argparse.SUPPRESS): | |
"""Initializes _FlagAction. | |
Args: | |
option_strings: See argparse.Action. | |
dest: Ignored. The flag is always defined with dest=argparse.SUPPRESS. | |
help: See argparse.Action. | |
metavar: See argparse.Action. | |
flag_instance: absl.flags.Flag, the absl flag instance. | |
default: Ignored. The flag always uses dest=argparse.SUPPRESS so it | |
doesn't affect the parsing result. | |
""" | |
del dest | |
self._flag_instance = flag_instance | |
super(_FlagAction, self).__init__( | |
option_strings=option_strings, | |
dest=argparse.SUPPRESS, | |
help=help, | |
metavar=metavar) | |
def __call__(self, parser, namespace, values, option_string=None): | |
"""See https://docs.python.org/3/library/argparse.html#action-classes.""" | |
self._flag_instance.parse(values) | |
self._flag_instance.using_default_value = False | |
class _BooleanFlagAction(argparse.Action): | |
"""Action class for Abseil boolean flags.""" | |
def __init__( | |
self, | |
option_strings, | |
dest, | |
help, # pylint: disable=redefined-builtin | |
metavar, | |
flag_instance, | |
default=argparse.SUPPRESS): | |
"""Initializes _BooleanFlagAction. | |
Args: | |
option_strings: See argparse.Action. | |
dest: Ignored. The flag is always defined with dest=argparse.SUPPRESS. | |
help: See argparse.Action. | |
metavar: See argparse.Action. | |
flag_instance: absl.flags.Flag, the absl flag instance. | |
default: Ignored. The flag always uses dest=argparse.SUPPRESS so it | |
doesn't affect the parsing result. | |
""" | |
del dest, default | |
self._flag_instance = flag_instance | |
flag_names = [self._flag_instance.name] | |
if self._flag_instance.short_name: | |
flag_names.append(self._flag_instance.short_name) | |
self._flag_names = frozenset(flag_names) | |
super(_BooleanFlagAction, self).__init__( | |
option_strings=option_strings, | |
dest=argparse.SUPPRESS, | |
nargs=0, # Does not accept values, only `--bool` or `--nobool`. | |
help=help, | |
metavar=metavar) | |
def __call__(self, parser, namespace, values, option_string=None): | |
"""See https://docs.python.org/3/library/argparse.html#action-classes.""" | |
if not isinstance(values, list) or values: | |
raise ValueError('values must be an empty list.') | |
if option_string.startswith('--'): | |
option = option_string[2:] | |
else: | |
option = option_string[1:] | |
if option in self._flag_names: | |
self._flag_instance.parse('true') | |
else: | |
if not option.startswith('no') or option[2:] not in self._flag_names: | |
raise ValueError('invalid option_string: ' + option_string) | |
self._flag_instance.parse('false') | |
self._flag_instance.using_default_value = False | |
class _HelpFullAction(argparse.Action): | |
"""Action class for --helpfull flag.""" | |
def __init__(self, option_strings, dest, default, help): # pylint: disable=redefined-builtin | |
"""Initializes _HelpFullAction. | |
Args: | |
option_strings: See argparse.Action. | |
dest: Ignored. The flag is always defined with dest=argparse.SUPPRESS. | |
default: Ignored. | |
help: See argparse.Action. | |
""" | |
del dest, default | |
super(_HelpFullAction, self).__init__( | |
option_strings=option_strings, | |
dest=argparse.SUPPRESS, | |
default=argparse.SUPPRESS, | |
nargs=0, | |
help=help) | |
def __call__(self, parser, namespace, values, option_string=None): | |
"""See https://docs.python.org/3/library/argparse.html#action-classes.""" | |
# This only prints flags when help is not argparse.SUPPRESS. | |
# It includes user defined argparse flags, as well as main module's | |
# key absl flags. Other absl flags use argparse.SUPPRESS, so they aren't | |
# printed here. | |
parser.print_help() | |
absl_flags = parser._inherited_absl_flags # pylint: disable=protected-access | |
if absl_flags is not None: | |
modules = sorted(absl_flags.flags_by_module_dict()) | |
main_module = sys.argv[0] | |
if main_module in modules: | |
# The main module flags are already printed in parser.print_help(). | |
modules.remove(main_module) | |
print(absl_flags._get_help_for_modules( # pylint: disable=protected-access | |
modules, prefix='', include_special_flags=True)) | |
parser.exit() | |
def _strip_undefok_args(undefok, args): | |
"""Returns a new list of args after removing flags in --undefok.""" | |
if undefok: | |
undefok_names = set(name.strip() for name in undefok.split(',')) | |
undefok_names |= set('no' + name for name in undefok_names) | |
# Remove undefok flags. | |
args = [arg for arg in args if not _is_undefok(arg, undefok_names)] | |
return args | |
def _is_undefok(arg, undefok_names): | |
"""Returns whether we can ignore arg based on a set of undefok flag names.""" | |
if not arg.startswith('-'): | |
return False | |
if arg.startswith('--'): | |
arg_without_dash = arg[2:] | |
else: | |
arg_without_dash = arg[1:] | |
if '=' in arg_without_dash: | |
name, _ = arg_without_dash.split('=', 1) | |
else: | |
name = arg_without_dash | |
if name in undefok_names: | |
return True | |
return False | |