461 lines
14 KiB
Python
461 lines
14 KiB
Python
|
# Copyright 2017 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.
|
||
|
|
|
||
|
|
"""Generic entry point for Abseil Python applications.
|
||
|
|
|
||
|
|
To use this module, define a 'main' function with a single 'argv' argument and
|
||
|
|
call app.run(main). For example:
|
||
|
|
|
||
|
|
def main(argv):
|
||
|
|
if len(argv) > 1:
|
||
|
|
raise app.UsageError('Too many command-line arguments.')
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
app.run(main)
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import absolute_import
|
||
|
|
from __future__ import division
|
||
|
|
from __future__ import print_function
|
||
|
|
|
||
|
|
import collections
|
||
|
|
import errno
|
||
|
|
import os
|
||
|
|
import pdb
|
||
|
|
import sys
|
||
|
|
import traceback
|
||
|
|
|
||
|
|
from absl import command_name
|
||
|
|
from absl import flags
|
||
|
|
from absl import logging
|
||
|
|
|
||
|
|
try:
|
||
|
|
import faulthandler
|
||
|
|
except ImportError:
|
||
|
|
faulthandler = None
|
||
|
|
|
||
|
|
FLAGS = flags.FLAGS
|
||
|
|
|
||
|
|
flags.DEFINE_boolean('run_with_pdb', False, 'Set to true for PDB debug mode')
|
||
|
|
flags.DEFINE_boolean('pdb_post_mortem', False,
|
||
|
|
'Set to true to handle uncaught exceptions with PDB '
|
||
|
|
'post mortem.')
|
||
|
|
flags.DEFINE_boolean('run_with_profiling', False,
|
||
|
|
'Set to true for profiling the script. '
|
||
|
|
'Execution will be slower, and the output format might '
|
||
|
|
'change over time.')
|
||
|
|
flags.DEFINE_string('profile_file', None,
|
||
|
|
'Dump profile information to a file (for python -m '
|
||
|
|
'pstats). Implies --run_with_profiling.')
|
||
|
|
flags.DEFINE_boolean('use_cprofile_for_profiling', True,
|
||
|
|
'Use cProfile instead of the profile module for '
|
||
|
|
'profiling. This has no effect unless '
|
||
|
|
'--run_with_profiling is set.')
|
||
|
|
flags.DEFINE_boolean('only_check_args', False,
|
||
|
|
'Set to true to validate args and exit.',
|
||
|
|
allow_hide_cpp=True)
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
# If main() exits via an abnormal exception, call into these
|
||
|
|
# handlers before exiting.
|
||
|
|
EXCEPTION_HANDLERS = []
|
||
|
|
|
||
|
|
|
||
|
|
class Error(Exception):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class UsageError(Error):
|
||
|
|
"""Exception raised when the arguments supplied by the user are invalid.
|
||
|
|
|
||
|
|
Raise this when the arguments supplied are invalid from the point of
|
||
|
|
view of the application. For example when two mutually exclusive
|
||
|
|
flags have been supplied or when there are not enough non-flag
|
||
|
|
arguments. It is distinct from flags.Error which covers the lower
|
||
|
|
level of parsing and validating individual flags.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, message, exitcode=1):
|
||
|
|
super(UsageError, self).__init__(message)
|
||
|
|
self.exitcode = exitcode
|
||
|
|
|
||
|
|
|
||
|
|
class HelpFlag(flags.BooleanFlag):
|
||
|
|
"""Special boolean flag that displays usage and raises SystemExit."""
|
||
|
|
NAME = 'help'
|
||
|
|
SHORT_NAME = '?'
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
super(HelpFlag, self).__init__(
|
||
|
|
self.NAME, False, 'show this help',
|
||
|
|
short_name=self.SHORT_NAME, allow_hide_cpp=True)
|
||
|
|
|
||
|
|
def parse(self, arg):
|
||
|
|
if arg:
|
||
|
|
usage(shorthelp=True, writeto_stdout=True)
|
||
|
|
# Advertise --helpfull on stdout, since usage() was on stdout.
|
||
|
|
print()
|
||
|
|
print('Try --helpfull to get a list of all flags.')
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
|
||
|
|
class HelpshortFlag(HelpFlag):
|
||
|
|
"""--helpshort is an alias for --help."""
|
||
|
|
NAME = 'helpshort'
|
||
|
|
SHORT_NAME = None
|
||
|
|
|
||
|
|
|
||
|
|
class HelpfullFlag(flags.BooleanFlag):
|
||
|
|
"""Display help for flags in the main module and all dependent modules."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
super(HelpfullFlag, self).__init__(
|
||
|
|
'helpfull', False, 'show full help', allow_hide_cpp=True)
|
||
|
|
|
||
|
|
def parse(self, arg):
|
||
|
|
if arg:
|
||
|
|
usage(writeto_stdout=True)
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
|
||
|
|
class HelpXMLFlag(flags.BooleanFlag):
|
||
|
|
"""Similar to HelpfullFlag, but generates output in XML format."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
super(HelpXMLFlag, self).__init__(
|
||
|
|
'helpxml', False, 'like --helpfull, but generates XML output',
|
||
|
|
allow_hide_cpp=True)
|
||
|
|
|
||
|
|
def parse(self, arg):
|
||
|
|
if arg:
|
||
|
|
flags.FLAGS.write_help_in_xml_format(sys.stdout)
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
|
||
|
|
def parse_flags_with_usage(args):
|
||
|
|
"""Tries to parse the flags, print usage, and exit if unparseable.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
args: [str], a non-empty list of the command line arguments including
|
||
|
|
program name.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
[str], a non-empty list of remaining command line arguments after parsing
|
||
|
|
flags, including program name.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
return FLAGS(args)
|
||
|
|
except flags.Error as error:
|
||
|
|
sys.stderr.write('FATAL Flags parsing error: %s\n' % error)
|
||
|
|
sys.stderr.write('Pass --helpshort or --helpfull to see help on flags.\n')
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
|
||
|
|
_define_help_flags_called = False
|
||
|
|
|
||
|
|
|
||
|
|
def define_help_flags():
|
||
|
|
"""Registers help flags. Idempotent."""
|
||
|
|
# Use a global to ensure idempotence.
|
||
|
|
global _define_help_flags_called
|
||
|
|
|
||
|
|
if not _define_help_flags_called:
|
||
|
|
flags.DEFINE_flag(HelpFlag())
|
||
|
|
flags.DEFINE_flag(HelpshortFlag()) # alias for --help
|
||
|
|
flags.DEFINE_flag(HelpfullFlag())
|
||
|
|
flags.DEFINE_flag(HelpXMLFlag())
|
||
|
|
_define_help_flags_called = True
|
||
|
|
|
||
|
|
|
||
|
|
def _register_and_parse_flags_with_usage(
|
||
|
|
argv=None,
|
||
|
|
flags_parser=parse_flags_with_usage,
|
||
|
|
):
|
||
|
|
"""Registers help flags, parses arguments and shows usage if appropriate.
|
||
|
|
|
||
|
|
This also calls sys.exit(0) if flag --only_check_args is True.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
argv: [str], a non-empty list of the command line arguments including
|
||
|
|
program name, sys.argv is used if None.
|
||
|
|
flags_parser: Callable[[List[Text]], Any], the function used to parse flags.
|
||
|
|
The return value of this function is passed to `main` untouched.
|
||
|
|
It must guarantee FLAGS is parsed after this function is called.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The return value of `flags_parser`. When using the default `flags_parser`,
|
||
|
|
it returns the following:
|
||
|
|
[str], a non-empty list of remaining command line arguments after parsing
|
||
|
|
flags, including program name.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
Error: Raised when flags_parser is called, but FLAGS is not parsed.
|
||
|
|
SystemError: Raised when it's called more than once.
|
||
|
|
"""
|
||
|
|
if _register_and_parse_flags_with_usage.done:
|
||
|
|
raise SystemError('Flag registration can be done only once.')
|
||
|
|
|
||
|
|
define_help_flags()
|
||
|
|
|
||
|
|
original_argv = sys.argv if argv is None else argv
|
||
|
|
args_to_main = flags_parser(original_argv)
|
||
|
|
if not FLAGS.is_parsed():
|
||
|
|
raise Error('FLAGS must be parsed after flags_parser is called.')
|
||
|
|
|
||
|
|
# Exit when told so.
|
||
|
|
if FLAGS.only_check_args:
|
||
|
|
sys.exit(0)
|
||
|
|
# Immediately after flags are parsed, bump verbosity to INFO if the flag has
|
||
|
|
# not been set.
|
||
|
|
if FLAGS['verbosity'].using_default_value:
|
||
|
|
FLAGS.verbosity = 0
|
||
|
|
_register_and_parse_flags_with_usage.done = True
|
||
|
|
|
||
|
|
return args_to_main
|
||
|
|
|
||
|
|
_register_and_parse_flags_with_usage.done = False
|
||
|
|
|
||
|
|
|
||
|
|
def _run_main(main, argv):
|
||
|
|
"""Calls main, optionally with pdb or profiler."""
|
||
|
|
if FLAGS.run_with_pdb:
|
||
|
|
sys.exit(pdb.runcall(main, argv))
|
||
|
|
elif FLAGS.run_with_profiling or FLAGS.profile_file:
|
||
|
|
# Avoid import overhead since most apps (including performance-sensitive
|
||
|
|
# ones) won't be run with profiling.
|
||
|
|
import atexit
|
||
|
|
if FLAGS.use_cprofile_for_profiling:
|
||
|
|
import cProfile as profile
|
||
|
|
else:
|
||
|
|
import profile
|
||
|
|
profiler = profile.Profile()
|
||
|
|
if FLAGS.profile_file:
|
||
|
|
atexit.register(profiler.dump_stats, FLAGS.profile_file)
|
||
|
|
else:
|
||
|
|
atexit.register(profiler.print_stats)
|
||
|
|
retval = profiler.runcall(main, argv)
|
||
|
|
sys.exit(retval)
|
||
|
|
else:
|
||
|
|
sys.exit(main(argv))
|
||
|
|
|
||
|
|
|
||
|
|
def _call_exception_handlers(exception):
|
||
|
|
"""Calls any installed exception handlers."""
|
||
|
|
for handler in EXCEPTION_HANDLERS:
|
||
|
|
try:
|
||
|
|
if handler.wants(exception):
|
||
|
|
handler.handle(exception)
|
||
|
|
except: # pylint: disable=bare-except
|
||
|
|
try:
|
||
|
|
# We don't want to stop for exceptions in the exception handlers but
|
||
|
|
# we shouldn't hide them either.
|
||
|
|
logging.error(traceback.format_exc())
|
||
|
|
except: # pylint: disable=bare-except
|
||
|
|
# In case even the logging statement fails, ignore.
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def run(
|
||
|
|
main,
|
||
|
|
argv=None,
|
||
|
|
flags_parser=parse_flags_with_usage,
|
||
|
|
):
|
||
|
|
"""Begins executing the program.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
main: The main function to execute. It takes an single argument "argv",
|
||
|
|
which is a list of command line arguments with parsed flags removed.
|
||
|
|
If it returns an integer, it is used as the process's exit code.
|
||
|
|
argv: A non-empty list of the command line arguments including program name,
|
||
|
|
sys.argv is used if None.
|
||
|
|
flags_parser: Callable[[List[Text]], Any], the function used to parse flags.
|
||
|
|
The return value of this function is passed to `main` untouched.
|
||
|
|
It must guarantee FLAGS is parsed after this function is called.
|
||
|
|
- Parses command line flags with the flag module.
|
||
|
|
- If there are any errors, prints usage().
|
||
|
|
- Calls main() with the remaining arguments.
|
||
|
|
- If main() raises a UsageError, prints usage and the error message.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
args = _run_init(
|
||
|
|
sys.argv if argv is None else argv,
|
||
|
|
flags_parser,
|
||
|
|
)
|
||
|
|
while _init_callbacks:
|
||
|
|
callback = _init_callbacks.popleft()
|
||
|
|
callback()
|
||
|
|
try:
|
||
|
|
_run_main(main, args)
|
||
|
|
except UsageError as error:
|
||
|
|
usage(shorthelp=True, detailed_error=error, exitcode=error.exitcode)
|
||
|
|
except:
|
||
|
|
if FLAGS.pdb_post_mortem:
|
||
|
|
traceback.print_exc()
|
||
|
|
pdb.post_mortem()
|
||
|
|
raise
|
||
|
|
except Exception as e:
|
||
|
|
_call_exception_handlers(e)
|
||
|
|
raise
|
||
|
|
|
||
|
|
# Callbacks which have been deferred until after _run_init has been called.
|
||
|
|
_init_callbacks = collections.deque()
|
||
|
|
|
||
|
|
|
||
|
|
def call_after_init(callback):
|
||
|
|
"""Calls the given callback only once ABSL has finished initialization.
|
||
|
|
|
||
|
|
If ABSL has already finished initialization when `call_after_init` is
|
||
|
|
called then the callback is executed immediately, otherwise `callback` is
|
||
|
|
stored to be executed after `app.run` has finished initializing (aka. just
|
||
|
|
before the main function is called).
|
||
|
|
|
||
|
|
If called after `app.run`, this is equivalent to calling `callback()` in the
|
||
|
|
caller thread. If called before `app.run`, callbacks are run sequentially (in
|
||
|
|
an undefined order) in the same thread as `app.run`.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
callback: a callable to be called once ABSL has finished initialization.
|
||
|
|
This may be immediate if initialization has already finished. It
|
||
|
|
takes no arguments and returns nothing.
|
||
|
|
"""
|
||
|
|
if _run_init.done:
|
||
|
|
callback()
|
||
|
|
else:
|
||
|
|
_init_callbacks.append(callback)
|
||
|
|
|
||
|
|
|
||
|
|
def _run_init(
|
||
|
|
argv,
|
||
|
|
flags_parser,
|
||
|
|
):
|
||
|
|
"""Does one-time initialization and re-parses flags on rerun."""
|
||
|
|
if _run_init.done:
|
||
|
|
return flags_parser(argv)
|
||
|
|
command_name.make_process_name_useful()
|
||
|
|
# Set up absl logging handler.
|
||
|
|
logging.use_absl_handler()
|
||
|
|
args = _register_and_parse_flags_with_usage(
|
||
|
|
argv=argv,
|
||
|
|
flags_parser=flags_parser,
|
||
|
|
)
|
||
|
|
if faulthandler:
|
||
|
|
try:
|
||
|
|
faulthandler.enable()
|
||
|
|
except Exception: # pylint: disable=broad-except
|
||
|
|
# Some tests verify stderr output very closely, so don't print anything.
|
||
|
|
# Disabled faulthandler is a low-impact error.
|
||
|
|
pass
|
||
|
|
_run_init.done = True
|
||
|
|
return args
|
||
|
|
|
||
|
|
|
||
|
|
_run_init.done = False
|
||
|
|
|
||
|
|
|
||
|
|
def usage(shorthelp=False, writeto_stdout=False, detailed_error=None,
|
||
|
|
exitcode=None):
|
||
|
|
"""Writes __main__'s docstring to stderr with some help text.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
shorthelp: bool, if True, prints only flags from the main module,
|
||
|
|
rather than all flags.
|
||
|
|
writeto_stdout: bool, if True, writes help message to stdout,
|
||
|
|
rather than to stderr.
|
||
|
|
detailed_error: str, additional detail about why usage info was presented.
|
||
|
|
exitcode: optional integer, if set, exits with this status code after
|
||
|
|
writing help.
|
||
|
|
"""
|
||
|
|
if writeto_stdout:
|
||
|
|
stdfile = sys.stdout
|
||
|
|
else:
|
||
|
|
stdfile = sys.stderr
|
||
|
|
|
||
|
|
doc = sys.modules['__main__'].__doc__
|
||
|
|
if not doc:
|
||
|
|
doc = '\nUSAGE: %s [flags]\n' % sys.argv[0]
|
||
|
|
doc = flags.text_wrap(doc, indent=' ', firstline_indent='')
|
||
|
|
else:
|
||
|
|
# Replace all '%s' with sys.argv[0], and all '%%' with '%'.
|
||
|
|
num_specifiers = doc.count('%') - 2 * doc.count('%%')
|
||
|
|
try:
|
||
|
|
doc %= (sys.argv[0],) * num_specifiers
|
||
|
|
except (OverflowError, TypeError, ValueError):
|
||
|
|
# Just display the docstring as-is.
|
||
|
|
pass
|
||
|
|
if shorthelp:
|
||
|
|
flag_str = FLAGS.main_module_help()
|
||
|
|
else:
|
||
|
|
flag_str = FLAGS.get_help()
|
||
|
|
try:
|
||
|
|
stdfile.write(doc)
|
||
|
|
if flag_str:
|
||
|
|
stdfile.write('\nflags:\n')
|
||
|
|
stdfile.write(flag_str)
|
||
|
|
stdfile.write('\n')
|
||
|
|
if detailed_error is not None:
|
||
|
|
stdfile.write('\n%s\n' % detailed_error)
|
||
|
|
except IOError as e:
|
||
|
|
# We avoid printing a huge backtrace if we get EPIPE, because
|
||
|
|
# "foo.par --help | less" is a frequent use case.
|
||
|
|
if e.errno != errno.EPIPE:
|
||
|
|
raise
|
||
|
|
if exitcode is not None:
|
||
|
|
sys.exit(exitcode)
|
||
|
|
|
||
|
|
|
||
|
|
class ExceptionHandler(object):
|
||
|
|
"""Base exception handler from which other may inherit."""
|
||
|
|
|
||
|
|
def wants(self, exc):
|
||
|
|
"""Returns whether this handler wants to handle the exception or not.
|
||
|
|
|
||
|
|
This base class returns True for all exceptions by default. Override in
|
||
|
|
subclass if it wants to be more selective.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
exc: Exception, the current exception.
|
||
|
|
"""
|
||
|
|
del exc # Unused.
|
||
|
|
return True
|
||
|
|
|
||
|
|
def handle(self, exc):
|
||
|
|
"""Do something with the current exception.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
exc: Exception, the current exception
|
||
|
|
|
||
|
|
This method must be overridden.
|
||
|
|
"""
|
||
|
|
raise NotImplementedError()
|
||
|
|
|
||
|
|
|
||
|
|
def install_exception_handler(handler):
|
||
|
|
"""Installs an exception handler.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
handler: ExceptionHandler, the exception handler to install.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
TypeError: Raised when the handler was not of the correct type.
|
||
|
|
|
||
|
|
All installed exception handlers will be called if main() exits via
|
||
|
|
an abnormal exception, i.e. not one of SystemExit, KeyboardInterrupt,
|
||
|
|
FlagsError or UsageError.
|
||
|
|
"""
|
||
|
|
if not isinstance(handler, ExceptionHandler):
|
||
|
|
raise TypeError('handler of type %s does not inherit from ExceptionHandler'
|
||
|
|
% type(handler))
|
||
|
|
EXCEPTION_HANDLERS.append(handler)
|