Create `app.call_after_init(callback)`.

This allows users to register callbacks which are guaranteed to run after
`app._run_init` has finished and before app.run's main function is invoked.
This is most useful to guarantee setup/init of something is performed when
a library is used under absl.

These may be useful to register callbacks as part of modules being imported
(which usually happens before `app.run`) or in library code that may run in
a context where ABSL's app is not guaranteed to be used by the end-user
(this would also assume that never running the callback is acceptable in that
case).

PiperOrigin-RevId: 205123575

2 files changed, 21 insertions, 3 deletions

diff --git a/absl/tests/app_test.py b/absl/tests/app_test.py
index 932343a..9fca67f 100644
--- a/absl/tests/app_test.py
+++ b/absl/tests/app_test.py
@@ -287,7 +287,7 @@ class FunctionalTests(absltest.TestCase):
         expected_stderr_substring=app_test_helper.__doc__)
     self.assertEqual(88, exitcode)
 
-  def test_exeption_handler(self):
+  def test_exception_handler(self):
     exception_handler_messages = (
         'MyExceptionHandler: first\nMyExceptionHandler: second\n')
     self.run_helper(
@@ -295,10 +295,16 @@ class FunctionalTests(absltest.TestCase):
         arguments=['--raise_exception'],
         expected_stdout_substring=exception_handler_messages)
 
-  def test_exeption_handler_not_called(self):
+  def test_exception_handler_not_called(self):
     _, _, stdout = self.run_helper(True)
     self.assertNotIn('MyExceptionHandler', stdout)
 
+  def test_print_init_callbacks(self):
+    _, stdout, _ = self.run_helper(
+        expect_success=True, arguments=['--print_init_callbacks'])
+    self.assertIn('before app.run', stdout)
+    self.assertIn('during real_main', stdout)
+
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/absl/tests/app_test_helper.py b/absl/tests/app_test_helper.py
index bed141b..cf2753e 100755
--- a/absl/tests/app_test_helper.py
+++ b/absl/tests/app_test_helper.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Helper script used by app_unittest.sh."""
+"""Helper script used by app_test.py."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -38,6 +38,8 @@ flags.DEFINE_integer(
     'usage_error_exitcode', None, 'The exitcode if app.UsageError if raised.')
 flags.DEFINE_string(
     'str_flag_with_unicode_args', u'thumb:\U0001F44D', u'smile:\U0001F604')
+flags.DEFINE_boolean('print_init_callbacks', False,
+                     'print init callbacks and exit')
 
 
 class MyException(Exception):
@@ -71,6 +73,12 @@ def real_main(argv):
     faulthandler._sigsegv()  # pylint: disable=protected-access
     sys.exit(1)  # Should not reach here.
 
+  if FLAGS.print_init_callbacks:
+    app.call_after_init(lambda: _callback_results.append('during real_main'))
+    for value in _callback_results:
+      print('callback: {}'.format(value))
+    sys.exit(0)
+
   # Ensure that we have a random C++ flag in flags.FLAGS; this shows
   # us that app.run() did the right thing in conjunction with C++ flags.
   helper_type = os.environ['APP_TEST_HELPER_TYPE']
@@ -103,6 +111,9 @@ def main(argv):
   real_main(argv)
 
 
+# Holds results from callbacks triggered by `app.run_after_init`.
+_callback_results = []
+
 if __name__ == '__main__':
   kwargs = {'main': main}
   main_function_name = os.environ.get('APP_TEST_CUSTOM_MAIN_FUNC', None)
@@ -112,6 +123,7 @@ if __name__ == '__main__':
   if custom_argv:
     kwargs['argv'] = custom_argv.split(' ')
 
+  app.call_after_init(lambda: _callback_results.append('before app.run'))
   app.install_exception_handler(MyExceptionHandler('first'))
   app.install_exception_handler(MyExceptionHandler('second'))
   app.run(**kwargs)