Allow $PYTHONBREAKPOINT to affect runcall and post_mortem debugging

PiperOrigin-RevId: 780556187

Author: Abseil Team

absl/BUILD

@@ -55,6 +55,7 @@ py_test(
         "//absl/testing:_bazelize_command",
         "//absl/testing:absltest",
         "//absl/testing:flagsaver",
+        "//absl/testing:parameterized",
     ],
 )
 

absl/app.py

@@ -27,6 +27,7 @@ def main(argv):
 
 import collections
 import errno
+import importlib
 import os
 import pdb
 import sys
@@ -44,10 +45,21 @@ def main(argv):
 
 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_pdb',
+    False,
+    'Set to true for debug mode. PDB is used by default; $PYTHONBREAKPOINT '
+    '(https://docs.python.org/3/using/cmdline.html#envvar-PYTHONBREAKPOINT) '
+    'can be used to specify a custom debugger.',
+)
+flags.DEFINE_boolean(
+    'pdb_post_mortem',
+    False,
+    'Set to true to handle uncaught exceptions with the post mortem debugger.'
+    'PDB is used by default; $PYTHONBREAKPOINT '
+    '(https://docs.python.org/3/using/cmdline.html#envvar-PYTHONBREAKPOINT) '
+    'can be used to specify a custom one.',
+)
 flags.DEFINE_alias('pdb', 'pdb_post_mortem')
 flags.DEFINE_boolean('run_with_profiling', False,
                      'Set to true for profiling the script. '
@@ -65,6 +77,45 @@ def main(argv):
                      allow_hide_cpp=True)
 
 
+def _get_debugger_module_with_function(function_name):
+  """Provides the `$PYTHONBREAKPOINT` module if it contains `function_name`.
+
+  Falls back to `pdb` otherwise.
+
+  Args:
+    function_name: The name of the function required.
+
+  Returns:
+    A debugger module providing `function_name`.
+  """
+  python_breakpoint = os.getenv('PYTHONBREAKPOINT')
+  # The special value '0' for `$PYTHONBREAKPOINT` means "do not use a debugger".
+  # We don't respect it (if the user explicitly asks to debug) but shouldn't try
+  # to import a module with this name.
+  if python_breakpoint and python_breakpoint != '0':
+    debugger_module_import = python_breakpoint.rsplit('.', 1)[0]
+    try:
+      debugger_module = importlib.import_module(debugger_module_import)
+    except ImportError:
+      logging.warning(
+          (
+              'Could not import $PYTHONBREAKPOINT debugger module %r, '
+              'falling back to PDB'
+          ),
+          debugger_module_import,
+      )
+    else:
+      if hasattr(debugger_module, function_name):
+        return debugger_module
+      logging.warning(
+          '$PYTHONBREAKPOINT debugger %r has no function %r, '
+          'falling back to PDB',
+          debugger_module_import,
+          function_name,
+      )
+  return pdb
+
+
 # If main() exits via an abnormal exception, call into these
 # handlers before exiting.
 EXCEPTION_HANDLERS = []
@@ -239,9 +290,9 @@ def _register_and_parse_flags_with_usage(
 
 
 def _run_main(main, argv):
-  """Calls main, optionally with pdb or profiler."""
+  """Calls main, optionally with a debugger or profiler."""
   if FLAGS.run_with_pdb:
-    sys.exit(pdb.runcall(main, argv))
+    sys.exit(_get_debugger_module_with_function('runcall').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.
@@ -331,7 +382,7 @@ def run(
         print()
         print(' *** Entering post-mortem debugging ***')
         print()
-        pdb.post_mortem()
+        _get_debugger_module_with_function('post_mortem').post_mortem()
       raise
   except Exception as e:
     _call_exception_handlers(e)

absl/tests/app_test.py

@@ -17,8 +17,10 @@
 import contextlib
 import copy
 import enum
+import importlib
 import io
 import os
+import pdb
 import re
 import subprocess
 import sys
@@ -30,6 +32,7 @@
 from absl.testing import _bazelize_command
 from absl.testing import absltest
 from absl.testing import flagsaver
+from absl.testing import parameterized
 from absl.tests import app_test_helper
 
 
@@ -118,6 +121,89 @@ def test_register_and_parse_flags_with_usage_exits_on_second_run(self):
       app._register_and_parse_flags_with_usage()
 
 
+class _MyDebuggerModule:
+  """Imitates some aspects of the `pdb` interface."""
+
+  def post_mortem(self):
+    pass
+
+  def runcall(self, *args, **kwargs):
+    del args, kwargs
+
+
+class _IncompleteDebuggerModule:
+  """Provides `post_mortem` but not `runcall`."""
+
+  def post_mortem(self):
+    pass
+
+
+def _pythonbreakpoint_variable(value):
+  """Returns a context manager setting the $PYTHONBREAKPOINT env. variable."""
+  new_environment = {} if value is None else {'PYTHONBREAKPOINT': value}
+  return mock.patch.object(os, 'environ', new_environment)
+
+
+class PythonBreakpointTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+
+    importlib_import_module = importlib.import_module
+
+    def my_import_module(module, *args, **kwargs):
+      if module == 'my.debugger':
+        return _MyDebuggerModule()
+      if module == 'incomplete.debugger':
+        return _IncompleteDebuggerModule()
+      if module == '0':
+        # `PYTHONBREAKPOINT=0` is a special value.
+        raise ValueError('Should not try to import module `0`')
+      return importlib_import_module(module, *args, **kwargs)
+
+    self._mock_import_module = self.enter_context(
+        mock.patch.object(importlib, 'import_module', my_import_module)
+    )
+
+  @parameterized.named_parameters(
+      ('no_variable', None),
+      ('empty_variable', ''),
+      ('explicit_pdb', 'pdb.set_trace'),
+      ('unimportable_module', 'unimportable.module.set_trace'),
+      ('opt_out', '0'),
+  )
+  def test_python_breakpoint_pdb(self, value):
+    with _pythonbreakpoint_variable(value):
+      self.assertIs(app._get_debugger_module_with_function('runcall'), pdb)
+      self.assertIs(app._get_debugger_module_with_function('post_mortem'), pdb)
+
+  def test_my_debugger(self):
+    with _pythonbreakpoint_variable('my.debugger.set_trace'):
+      debugger_with_runcall = app._get_debugger_module_with_function('runcall')
+      self.assertIsInstance(debugger_with_runcall, _MyDebuggerModule)
+      self.assertTrue(hasattr(debugger_with_runcall, 'runcall'))
+
+      debugger_with_post_mortem = app._get_debugger_module_with_function(
+          'post_mortem'
+      )
+      self.assertIsInstance(debugger_with_post_mortem, _MyDebuggerModule)
+      self.assertTrue(hasattr(debugger_with_post_mortem, 'post_mortem'))
+
+  def test_incomplete_debugger(self):
+    with _pythonbreakpoint_variable('incomplete.debugger.set_trace'):
+      debugger_with_runcall = app._get_debugger_module_with_function('runcall')
+      self.assertIs(debugger_with_runcall, pdb)
+      self.assertTrue(hasattr(debugger_with_runcall, 'runcall'))
+
+      debugger_with_post_mortem = app._get_debugger_module_with_function(
+          'post_mortem'
+      )
+      self.assertIsInstance(
+          debugger_with_post_mortem, _IncompleteDebuggerModule
+      )
+      self.assertTrue(hasattr(debugger_with_post_mortem, 'post_mortem'))
+
+
 class FunctionalTests(absltest.TestCase):
   """Functional tests that use runs app_test_helper."""