Add support for stacklevel in the logging module as in the standard library.

Abseil Team · copybara-github · commit ae0a316c3dc1 · 2024-12-10T13:17:31.000-08:00
PiperOrigin-RevId: 704822441
diff --git a/absl/logging/__init__.py b/absl/logging/__init__.py
@@ -1118,10 +1118,11 @@ def findCaller(self, stack_info=False, stacklevel=1):
 
     Args:
       stack_info: bool, when True, include the stack trace as a fourth item
-          returned.  On Python 3 there are always four items returned - the
-          fourth will be None when this is False.  On Python 2 the stdlib
-          base class API only returns three items.  We do the same when this
-          new parameter is unspecified or False for compatibility.
+        returned.  On Python 3 there are always four items returned - the fourth
+        will be None when this is False.  On Python 2 the stdlib base class API
+        only returns three items.  We do the same when this new parameter is
+        unspecified or False for compatibility.
+      stacklevel: int, if greater than 1, that number of frames will be skipped.
 
     Returns:
       (filename, lineno, methodname[, sinfo]) of the calling method.
@@ -1130,22 +1131,36 @@ def findCaller(self, stack_info=False, stacklevel=1):
     # Use sys._getframe(2) instead of logging.currentframe(), it's slightly
     # faster because there is one less frame to traverse.
     frame = sys._getframe(2)  # pylint: disable=protected-access
+    frame_to_return = None
 
     while frame:
       code = frame.f_code
       if (_LOGGING_FILE_PREFIX not in code.co_filename and
           (code.co_filename, code.co_name,
            code.co_firstlineno) not in f_to_skip and
           (code.co_filename, code.co_name) not in f_to_skip):
-        sinfo = None
-        if stack_info:
-          out = io.StringIO()
-          out.write('Stack (most recent call last):\n')
-          traceback.print_stack(frame, file=out)
-          sinfo = out.getvalue().rstrip('\n')
-        return (code.co_filename, frame.f_lineno, code.co_name, sinfo)
+        frame_to_return = frame
+        stacklevel -= 1
+        if stacklevel <= 0:
+          break
       frame = frame.f_back
 
+    if frame_to_return is not None:
+      sinfo = None
+      if stack_info:
+        out = io.StringIO()
+        out.write('Stack (most recent call last):\n')
+        traceback.print_stack(frame, file=out)
+        sinfo = out.getvalue().rstrip('\n')
+      return (
+          frame_to_return.f_code.co_filename,
+          frame_to_return.f_lineno,
+          frame_to_return.f_code.co_name,
+          sinfo,
+      )
+
+    return None
+
   def critical(self, msg, *args, **kwargs):
     """Logs ``msg % args`` with severity ``CRITICAL``."""
     self.log(logging.CRITICAL, msg, *args, **kwargs)
diff --git a/absl/logging/tests/logging_test.py b/absl/logging/tests/logging_test.py
@@ -593,6 +593,76 @@ def test_logger_cannot_be_disabled(self):
       self.logger.handle(record)
     mock_call_handlers.assert_called_once()
 
+  def test_find_caller_respects_stacklevel(self):
+    self.set_up_mock_frames()
+    self.logger.register_frame_to_skip('myfile.py', 'Method1')
+    # Frame 0 is filtered out due to being in the logging file.
+    # Frame 1 is filtered out due to being registered to be skipped.
+    # Frame 2 is the topmost unfiltered frame on the stack.
+    self.assertEqual(
+        self.logger.findCaller(stacklevel=1),
+        ('myfile.py', 125, 'Method2', None),
+    )
+    self.assertEqual(
+        self.logger.findCaller(stacklevel=2),
+        ('myfile.py', 125, 'Method3', None),
+    )
+    self.assertEqual(
+        self.logger.findCaller(stacklevel=3),
+        ('myfile.py', 249, 'Method2', None),
+    )
+
+  def test_find_caller_handles_skipped_bottom_of_stack(self):
+    self.set_up_mock_frames()
+    self.logger.register_frame_to_skip('myfile.py', 'Method2')
+    # Frame 0 is filtered out due to being in the logging file.
+    # Frame 1 is unfiltered.
+    # Frame 2 is filtered due to regiser_frame_to_skip.
+    # Frame 3 is unfiltered.
+    # Frame 4 is filtered due to regiser_frame_to_skip.
+    self.assertEqual(
+        self.logger.findCaller(stacklevel=1),
+        ('myfile.py', 125, 'Method1', None),
+    )
+    self.assertEqual(
+        self.logger.findCaller(stacklevel=2),
+        ('myfile.py', 125, 'Method3', None),
+    )
+    # 3 exceeds the unfiltered stack depth, so it returns the bottom unfiltered frame.
+    self.assertEqual(
+        self.logger.findCaller(stacklevel=3),
+        ('myfile.py', 125, 'Method3', None),
+    )
+
+  def test_log_respects_stacklevel(self):
+    self.set_up_mock_frames()
+    original_logger_class = std_logging.getLoggerClass()
+    std_logging.setLoggerClass(logging.ABSLLogger)
+    absl_logger = std_logging.getLogger('absl')
+    with self.assertLogs() as cm:
+      absl_logger.info('lol', stacklevel=2)
+    std_logging.setLoggerClass(original_logger_class)
+    self.assertLen(cm.records, 1)
+    self.assertEqual(cm.records[0].lineno, 125)
+    self.assertEqual(cm.records[0].funcName, 'Method2')
+    self.assertEqual(cm.records[0].filename, 'myfile.py')
+
+  def test_handles_negative_stacklevel(self):
+    """If stacklevel is negative, return the top unfiltered frame."""
+    self.set_up_mock_frames()
+    self.assertEqual(
+        self.logger.findCaller(stacklevel=-1),
+        ('myfile.py', 125, 'Method1', None),
+    )
+
+  def test_handles_excess_stacklevel(self):
+    """If stacklevel exceeds the stack depth, return the bottom frame."""
+    self.set_up_mock_frames()
+    self.assertEqual(
+        self.logger.findCaller(stacklevel=1000),
+        ('myfile.py', 249, 'Method2', None),
+    )
+
 
 class ABSLLogPrefixTest(parameterized.TestCase):
 
