You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

460 lines
14 KiB

4 years ago
  1. # Copyright 2017 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. """Generic entry point for Abseil Python applications.
  15. To use this module, define a 'main' function with a single 'argv' argument and
  16. call app.run(main). For example:
  17. def main(argv):
  18. if len(argv) > 1:
  19. raise app.UsageError('Too many command-line arguments.')
  20. if __name__ == '__main__':
  21. app.run(main)
  22. """
  23. from __future__ import absolute_import
  24. from __future__ import division
  25. from __future__ import print_function
  26. import collections
  27. import errno
  28. import os
  29. import pdb
  30. import sys
  31. import traceback
  32. from absl import command_name
  33. from absl import flags
  34. from absl import logging
  35. try:
  36. import faulthandler
  37. except ImportError:
  38. faulthandler = None
  39. FLAGS = flags.FLAGS
  40. flags.DEFINE_boolean('run_with_pdb', False, 'Set to true for PDB debug mode')
  41. flags.DEFINE_boolean('pdb_post_mortem', False,
  42. 'Set to true to handle uncaught exceptions with PDB '
  43. 'post mortem.')
  44. flags.DEFINE_boolean('run_with_profiling', False,
  45. 'Set to true for profiling the script. '
  46. 'Execution will be slower, and the output format might '
  47. 'change over time.')
  48. flags.DEFINE_string('profile_file', None,
  49. 'Dump profile information to a file (for python -m '
  50. 'pstats). Implies --run_with_profiling.')
  51. flags.DEFINE_boolean('use_cprofile_for_profiling', True,
  52. 'Use cProfile instead of the profile module for '
  53. 'profiling. This has no effect unless '
  54. '--run_with_profiling is set.')
  55. flags.DEFINE_boolean('only_check_args', False,
  56. 'Set to true to validate args and exit.',
  57. allow_hide_cpp=True)
  58. # If main() exits via an abnormal exception, call into these
  59. # handlers before exiting.
  60. EXCEPTION_HANDLERS = []
  61. class Error(Exception):
  62. pass
  63. class UsageError(Error):
  64. """Exception raised when the arguments supplied by the user are invalid.
  65. Raise this when the arguments supplied are invalid from the point of
  66. view of the application. For example when two mutually exclusive
  67. flags have been supplied or when there are not enough non-flag
  68. arguments. It is distinct from flags.Error which covers the lower
  69. level of parsing and validating individual flags.
  70. """
  71. def __init__(self, message, exitcode=1):
  72. super(UsageError, self).__init__(message)
  73. self.exitcode = exitcode
  74. class HelpFlag(flags.BooleanFlag):
  75. """Special boolean flag that displays usage and raises SystemExit."""
  76. NAME = 'help'
  77. SHORT_NAME = '?'
  78. def __init__(self):
  79. super(HelpFlag, self).__init__(
  80. self.NAME, False, 'show this help',
  81. short_name=self.SHORT_NAME, allow_hide_cpp=True)
  82. def parse(self, arg):
  83. if arg:
  84. usage(shorthelp=True, writeto_stdout=True)
  85. # Advertise --helpfull on stdout, since usage() was on stdout.
  86. print()
  87. print('Try --helpfull to get a list of all flags.')
  88. sys.exit(1)
  89. class HelpshortFlag(HelpFlag):
  90. """--helpshort is an alias for --help."""
  91. NAME = 'helpshort'
  92. SHORT_NAME = None
  93. class HelpfullFlag(flags.BooleanFlag):
  94. """Display help for flags in the main module and all dependent modules."""
  95. def __init__(self):
  96. super(HelpfullFlag, self).__init__(
  97. 'helpfull', False, 'show full help', allow_hide_cpp=True)
  98. def parse(self, arg):
  99. if arg:
  100. usage(writeto_stdout=True)
  101. sys.exit(1)
  102. class HelpXMLFlag(flags.BooleanFlag):
  103. """Similar to HelpfullFlag, but generates output in XML format."""
  104. def __init__(self):
  105. super(HelpXMLFlag, self).__init__(
  106. 'helpxml', False, 'like --helpfull, but generates XML output',
  107. allow_hide_cpp=True)
  108. def parse(self, arg):
  109. if arg:
  110. flags.FLAGS.write_help_in_xml_format(sys.stdout)
  111. sys.exit(1)
  112. def parse_flags_with_usage(args):
  113. """Tries to parse the flags, print usage, and exit if unparseable.
  114. Args:
  115. args: [str], a non-empty list of the command line arguments including
  116. program name.
  117. Returns:
  118. [str], a non-empty list of remaining command line arguments after parsing
  119. flags, including program name.
  120. """
  121. try:
  122. return FLAGS(args)
  123. except flags.Error as error:
  124. sys.stderr.write('FATAL Flags parsing error: %s\n' % error)
  125. sys.stderr.write('Pass --helpshort or --helpfull to see help on flags.\n')
  126. sys.exit(1)
  127. _define_help_flags_called = False
  128. def define_help_flags():
  129. """Registers help flags. Idempotent."""
  130. # Use a global to ensure idempotence.
  131. global _define_help_flags_called
  132. if not _define_help_flags_called:
  133. flags.DEFINE_flag(HelpFlag())
  134. flags.DEFINE_flag(HelpshortFlag()) # alias for --help
  135. flags.DEFINE_flag(HelpfullFlag())
  136. flags.DEFINE_flag(HelpXMLFlag())
  137. _define_help_flags_called = True
  138. def _register_and_parse_flags_with_usage(
  139. argv=None,
  140. flags_parser=parse_flags_with_usage,
  141. ):
  142. """Registers help flags, parses arguments and shows usage if appropriate.
  143. This also calls sys.exit(0) if flag --only_check_args is True.
  144. Args:
  145. argv: [str], a non-empty list of the command line arguments including
  146. program name, sys.argv is used if None.
  147. flags_parser: Callable[[List[Text]], Any], the function used to parse flags.
  148. The return value of this function is passed to `main` untouched.
  149. It must guarantee FLAGS is parsed after this function is called.
  150. Returns:
  151. The return value of `flags_parser`. When using the default `flags_parser`,
  152. it returns the following:
  153. [str], a non-empty list of remaining command line arguments after parsing
  154. flags, including program name.
  155. Raises:
  156. Error: Raised when flags_parser is called, but FLAGS is not parsed.
  157. SystemError: Raised when it's called more than once.
  158. """
  159. if _register_and_parse_flags_with_usage.done:
  160. raise SystemError('Flag registration can be done only once.')
  161. define_help_flags()
  162. original_argv = sys.argv if argv is None else argv
  163. args_to_main = flags_parser(original_argv)
  164. if not FLAGS.is_parsed():
  165. raise Error('FLAGS must be parsed after flags_parser is called.')
  166. # Exit when told so.
  167. if FLAGS.only_check_args:
  168. sys.exit(0)
  169. # Immediately after flags are parsed, bump verbosity to INFO if the flag has
  170. # not been set.
  171. if FLAGS['verbosity'].using_default_value:
  172. FLAGS.verbosity = 0
  173. _register_and_parse_flags_with_usage.done = True
  174. return args_to_main
  175. _register_and_parse_flags_with_usage.done = False
  176. def _run_main(main, argv):
  177. """Calls main, optionally with pdb or profiler."""
  178. if FLAGS.run_with_pdb:
  179. sys.exit(pdb.runcall(main, argv))
  180. elif FLAGS.run_with_profiling or FLAGS.profile_file:
  181. # Avoid import overhead since most apps (including performance-sensitive
  182. # ones) won't be run with profiling.
  183. import atexit
  184. if FLAGS.use_cprofile_for_profiling:
  185. import cProfile as profile
  186. else:
  187. import profile
  188. profiler = profile.Profile()
  189. if FLAGS.profile_file:
  190. atexit.register(profiler.dump_stats, FLAGS.profile_file)
  191. else:
  192. atexit.register(profiler.print_stats)
  193. retval = profiler.runcall(main, argv)
  194. sys.exit(retval)
  195. else:
  196. sys.exit(main(argv))
  197. def _call_exception_handlers(exception):
  198. """Calls any installed exception handlers."""
  199. for handler in EXCEPTION_HANDLERS:
  200. try:
  201. if handler.wants(exception):
  202. handler.handle(exception)
  203. except: # pylint: disable=bare-except
  204. try:
  205. # We don't want to stop for exceptions in the exception handlers but
  206. # we shouldn't hide them either.
  207. logging.error(traceback.format_exc())
  208. except: # pylint: disable=bare-except
  209. # In case even the logging statement fails, ignore.
  210. pass
  211. def run(
  212. main,
  213. argv=None,
  214. flags_parser=parse_flags_with_usage,
  215. ):
  216. """Begins executing the program.
  217. Args:
  218. main: The main function to execute. It takes an single argument "argv",
  219. which is a list of command line arguments with parsed flags removed.
  220. If it returns an integer, it is used as the process's exit code.
  221. argv: A non-empty list of the command line arguments including program name,
  222. sys.argv is used if None.
  223. flags_parser: Callable[[List[Text]], Any], the function used to parse flags.
  224. The return value of this function is passed to `main` untouched.
  225. It must guarantee FLAGS is parsed after this function is called.
  226. - Parses command line flags with the flag module.
  227. - If there are any errors, prints usage().
  228. - Calls main() with the remaining arguments.
  229. - If main() raises a UsageError, prints usage and the error message.
  230. """
  231. try:
  232. args = _run_init(
  233. sys.argv if argv is None else argv,
  234. flags_parser,
  235. )
  236. while _init_callbacks:
  237. callback = _init_callbacks.popleft()
  238. callback()
  239. try:
  240. _run_main(main, args)
  241. except UsageError as error:
  242. usage(shorthelp=True, detailed_error=error, exitcode=error.exitcode)
  243. except:
  244. if FLAGS.pdb_post_mortem:
  245. traceback.print_exc()
  246. pdb.post_mortem()
  247. raise
  248. except Exception as e:
  249. _call_exception_handlers(e)
  250. raise
  251. # Callbacks which have been deferred until after _run_init has been called.
  252. _init_callbacks = collections.deque()
  253. def call_after_init(callback):
  254. """Calls the given callback only once ABSL has finished initialization.
  255. If ABSL has already finished initialization when `call_after_init` is
  256. called then the callback is executed immediately, otherwise `callback` is
  257. stored to be executed after `app.run` has finished initializing (aka. just
  258. before the main function is called).
  259. If called after `app.run`, this is equivalent to calling `callback()` in the
  260. caller thread. If called before `app.run`, callbacks are run sequentially (in
  261. an undefined order) in the same thread as `app.run`.
  262. Args:
  263. callback: a callable to be called once ABSL has finished initialization.
  264. This may be immediate if initialization has already finished. It
  265. takes no arguments and returns nothing.
  266. """
  267. if _run_init.done:
  268. callback()
  269. else:
  270. _init_callbacks.append(callback)
  271. def _run_init(
  272. argv,
  273. flags_parser,
  274. ):
  275. """Does one-time initialization and re-parses flags on rerun."""
  276. if _run_init.done:
  277. return flags_parser(argv)
  278. command_name.make_process_name_useful()
  279. # Set up absl logging handler.
  280. logging.use_absl_handler()
  281. args = _register_and_parse_flags_with_usage(
  282. argv=argv,
  283. flags_parser=flags_parser,
  284. )
  285. if faulthandler:
  286. try:
  287. faulthandler.enable()
  288. except Exception: # pylint: disable=broad-except
  289. # Some tests verify stderr output very closely, so don't print anything.
  290. # Disabled faulthandler is a low-impact error.
  291. pass
  292. _run_init.done = True
  293. return args
  294. _run_init.done = False
  295. def usage(shorthelp=False, writeto_stdout=False, detailed_error=None,
  296. exitcode=None):
  297. """Writes __main__'s docstring to stderr with some help text.
  298. Args:
  299. shorthelp: bool, if True, prints only flags from the main module,
  300. rather than all flags.
  301. writeto_stdout: bool, if True, writes help message to stdout,
  302. rather than to stderr.
  303. detailed_error: str, additional detail about why usage info was presented.
  304. exitcode: optional integer, if set, exits with this status code after
  305. writing help.
  306. """
  307. if writeto_stdout:
  308. stdfile = sys.stdout
  309. else:
  310. stdfile = sys.stderr
  311. doc = sys.modules['__main__'].__doc__
  312. if not doc:
  313. doc = '\nUSAGE: %s [flags]\n' % sys.argv[0]
  314. doc = flags.text_wrap(doc, indent=' ', firstline_indent='')
  315. else:
  316. # Replace all '%s' with sys.argv[0], and all '%%' with '%'.
  317. num_specifiers = doc.count('%') - 2 * doc.count('%%')
  318. try:
  319. doc %= (sys.argv[0],) * num_specifiers
  320. except (OverflowError, TypeError, ValueError):
  321. # Just display the docstring as-is.
  322. pass
  323. if shorthelp:
  324. flag_str = FLAGS.main_module_help()
  325. else:
  326. flag_str = FLAGS.get_help()
  327. try:
  328. stdfile.write(doc)
  329. if flag_str:
  330. stdfile.write('\nflags:\n')
  331. stdfile.write(flag_str)
  332. stdfile.write('\n')
  333. if detailed_error is not None:
  334. stdfile.write('\n%s\n' % detailed_error)
  335. except IOError as e:
  336. # We avoid printing a huge backtrace if we get EPIPE, because
  337. # "foo.par --help | less" is a frequent use case.
  338. if e.errno != errno.EPIPE:
  339. raise
  340. if exitcode is not None:
  341. sys.exit(exitcode)
  342. class ExceptionHandler(object):
  343. """Base exception handler from which other may inherit."""
  344. def wants(self, exc):
  345. """Returns whether this handler wants to handle the exception or not.
  346. This base class returns True for all exceptions by default. Override in
  347. subclass if it wants to be more selective.
  348. Args:
  349. exc: Exception, the current exception.
  350. """
  351. del exc # Unused.
  352. return True
  353. def handle(self, exc):
  354. """Do something with the current exception.
  355. Args:
  356. exc: Exception, the current exception
  357. This method must be overridden.
  358. """
  359. raise NotImplementedError()
  360. def install_exception_handler(handler):
  361. """Installs an exception handler.
  362. Args:
  363. handler: ExceptionHandler, the exception handler to install.
  364. Raises:
  365. TypeError: Raised when the handler was not of the correct type.
  366. All installed exception handlers will be called if main() exits via
  367. an abnormal exception, i.e. not one of SystemExit, KeyboardInterrupt,
  368. FlagsError or UsageError.
  369. """
  370. if not isinstance(handler, ExceptionHandler):
  371. raise TypeError('handler of type %s does not inherit from ExceptionHandler'
  372. % type(handler))
  373. EXCEPTION_HANDLERS.append(handler)