bpo-31299: Make it possible to filter out frames from tracebacks

Doc/library/traceback.rst

@@ -212,7 +212,7 @@ The module also defines the following classes:
 :class:`TracebackException` objects are created from actual exceptions to
 capture data for later printing in a lightweight fashion.
 
-.. class:: TracebackException(exc_type, exc_value, exc_traceback, *, limit=None, lookup_lines=True, capture_locals=False, compact=False)
+.. class:: TracebackException(exc_type, exc_value, exc_traceback, *, limit=None, lookup_lines=True, capture_locals=False, compact=False, stack_summary_cls=None)
 
    Capture an exception for later rendering. *limit*, *lookup_lines* and
    *capture_locals* are as for the :class:`StackSummary` class.
@@ -222,6 +222,10 @@ capture data for later printing in a lightweight fashion.
    ``__context__`` field is calculated only if ``__cause__`` is ``None`` and
    ``__suppress_context__`` is false.
 
+   If *stack_summary_cls* is not ``None``, it is a class to be used instead of
+   the default :class:`~traceback.StackSummary` to format the stack (typically
+   a subclass that overrides :meth:`~traceback.StackSummary.format_frame`).
+
    Note that when locals are captured, they are also shown in the traceback.
 
    .. attribute:: __cause__
@@ -309,6 +313,8 @@ capture data for later printing in a lightweight fashion.
    .. versionchanged:: 3.10
       Added the *compact* parameter.
 
+   .. versionchanged:: 3.11
+      Added the *stack_summary_cls* paramter.
 
 :class:`StackSummary` Objects
 -----------------------------
@@ -357,7 +363,8 @@ capture data for later printing in a lightweight fashion.
 
       Returns a string for printing one of the frames involved in the stack.
       This method gets called for each frame object to be printed in the
-      :class:`StackSummary`.
+      :class:`StackSummary`. If it returns ``None``, the frame is omitted
+      from the output.
 
       .. versionadded:: 3.11
 

Lib/test/test_traceback.py

@@ -1443,6 +1443,34 @@ def some_inner():
             s.format(),
             [f'{__file__}:{some_inner.__code__.co_firstlineno + 1}'])
 
+    def test_dropping_frames(self):
+        def f():
+            1/0
+
+        def g():
+            try:
+                f()
+            except:
+                return sys.exc_info()
+
+        exc_info = g()
+
+        class Skip_G(traceback.StackSummary):
+            def format_frame(self, frame):
+                if frame.name == 'g':
+                    return None
+                return super().format_frame(frame)
+
+        stack = Skip_G.extract(
+            traceback.walk_tb(exc_info[2])).format()
+
+        self.assertEqual(len(stack), 1)
+        lno = f.__code__.co_firstlineno + 1
+        self.assertEqual(
+            stack[0],
+            f'  File "{__file__}", line {lno}, in f\n    1/0\n'
+        )
+
 
 class TestTracebackException(unittest.TestCase):
 
@@ -1752,6 +1780,184 @@ def f():
              ''])
 
 
+class TestTracebackException_CustomStackSummary(unittest.TestCase):
+    def _get_output(self, *exc_info, stack_summary_cls=None):
+        output = StringIO()
+        traceback.TracebackException(
+            *exc_info, stack_summary_cls=stack_summary_cls,
+        ).print(file=output)
+        return output.getvalue().split('\n')
+
+    class MyStackSummary(traceback.StackSummary):
+        def format_frame(self, frame):
+            return f'{frame.filename}:{frame.lineno}\n'
+
+    class SkipG(traceback.StackSummary):
+        def format_frame(self, frame):
+            if frame.name == 'g':
+                return None
+            return super().format_frame(frame)
+
+    def test_custom_stack_summary(self):
+        def f():
+            1/0
+
+        def g():
+            try:
+                f()
+            except:
+                return sys.exc_info()
+
+        exc_info = g()
+        stack = self._get_output(
+            *exc_info,
+            stack_summary_cls=self.MyStackSummary)
+        self.assertEqual(
+            stack,
+            ['Traceback (most recent call last):',
+             f'{__file__}:{g.__code__.co_firstlineno+2}',
+             f'{__file__}:{f.__code__.co_firstlineno+1}',
+             'ZeroDivisionError: division by zero',
+             ''])
+
+    def test_custom_stack_summary_with_context(self):
+        def f():
+            try:
+                1/0
+            except ZeroDivisionError as e:
+                raise ValueError('bad value')
+
+        def g():
+            try:
+                f()
+            except:
+                return sys.exc_info()
+
+        exc_info = g()
+        stack = self._get_output(
+            *exc_info,
+            stack_summary_cls=self.MyStackSummary)
+        self.assertEqual(
+            stack,
+            ['Traceback (most recent call last):',
+             f'{__file__}:{f.__code__.co_firstlineno+2}',
+             'ZeroDivisionError: division by zero',
+             '',
+             context_message.replace('\n', ''),
+             '',
+             'Traceback (most recent call last):',
+             f'{__file__}:{g.__code__.co_firstlineno+2}',
+             f'{__file__}:{f.__code__.co_firstlineno+4}',
+             'ValueError: bad value',
+             ''])
+
+    def test_custom_stack_summary_with_cause(self):
+        def f():
+            try:
+                1/0
+            except ZeroDivisionError as e:
+                raise ValueError('bad value') from e
+
+        def g():
+            try:
+                f()
+            except:
+                return sys.exc_info()
+
+        exc_info = g()
+        stack = self._get_output(
+            *exc_info,
+            stack_summary_cls=self.MyStackSummary)
+        self.assertEqual(
+            stack,
+            ['Traceback (most recent call last):',
+             f'{__file__}:{f.__code__.co_firstlineno+2}',
+             'ZeroDivisionError: division by zero',
+             '',
+             cause_message.replace('\n', ''),
+             '',
+             'Traceback (most recent call last):',
+             f'{__file__}:{g.__code__.co_firstlineno+2}',
+             f'{__file__}:{f.__code__.co_firstlineno+4}',
+             'ValueError: bad value',
+             ''])
+
+    @requires_debug_ranges
+    def test_dropping_frames(self):
+        def f():
+            1/0
+
+        def g():
+            try:
+                f()
+            except:
+                return sys.exc_info()
+
+        exc_info = g()
+        full = self._get_output(*exc_info)
+        skipped = self._get_output(
+            *exc_info,
+            stack_summary_cls=self.SkipG)
+
+        for l in skipped:
+            full.remove(l)
+        # Only the lines for g's frame should remain:
+        self.assertEqual(len(full), 3)
+        lno = g.__code__.co_firstlineno + 2
+        self.assertEqual(
+            full,
+            [f'  File "{__file__}", line {lno}, in g',
+            '    f()',
+            '    ^^^'])
+
+    def test_dropping_frames_recursion_limit_msg1(self):
+        # recursion at bottom of the stack
+        def g():
+            g()
+
+        def h():
+            g()
+
+        try:
+            h()
+        except:
+            exc_info = sys.exc_info()
+
+        full = self._get_output(*exc_info)
+        skipped = self._get_output(
+            *exc_info,
+            stack_summary_cls=self.SkipG)
+
+        rep_txt_regex = 'Previous line repeated (\\d+) more times'
+        self.assertRegex(''.join(full), rep_txt_regex)
+        self.assertNotRegex(''.join(skipped), rep_txt_regex)
+
+    def test_dropping_frames_recursion_limit_msg2(self):
+        # recursion in the middle of the stack
+        def f():
+            1/0
+
+        def g(i):
+            if i < 10:
+                g(i+1)
+            else:
+                f()
+
+        try:
+            g(0)
+        except:
+            exc_info = sys.exc_info()
+
+        full = self._get_output(*exc_info)
+        skipped = self._get_output(
+            *exc_info,
+            stack_summary_cls=self.SkipG)
+
+        rep_txt_regex = 'Previous line repeated (\\d+) more times'
+        self.assertRegex(''.join(full), rep_txt_regex)
+        self.assertNotRegex(''.join(skipped), rep_txt_regex)
+
+
 class MiscTest(unittest.TestCase):
 
     def test_all(self):

Lib/traceback.py

@@ -453,7 +453,8 @@ def format_frame(self, frame):
         """Format the lines for a single frame.
 
         Returns a string representing one frame involved in the stack. This
-        gets called for every frame to be printed in the stack summary.
+        gets called for every frame to be printed in the stack summary. If
+        it returns ``None``, the frame is omitted from the output.
         """
         row = []
         row.append('  File "{}", line {}, in {}\n'.format(
@@ -509,23 +510,25 @@ def format(self):
         last_name = None
         count = 0
         for frame in self:
-            if (last_file is None or last_file != frame.filename or
-                last_line is None or last_line != frame.lineno or
-                last_name is None or last_name != frame.name):
+            formatted_frame = self.format_frame(frame)
+            if formatted_frame is not None:
+                if (last_file is None or last_file != frame.filename or
+                    last_line is None or last_line != frame.lineno or
+                    last_name is None or last_name != frame.name):
+                    if count > _RECURSIVE_CUTOFF:
+                        count -= _RECURSIVE_CUTOFF
+                        result.append(
+                            f'  [Previous line repeated {count} more '
+                            f'time{"s" if count > 1 else ""}]\n'
+                        )
+                    last_file = frame.filename
+                    last_line = frame.lineno
+                    last_name = frame.name
+                    count = 0
+                count += 1
                 if count > _RECURSIVE_CUTOFF:
-                    count -= _RECURSIVE_CUTOFF
-                    result.append(
-                        f'  [Previous line repeated {count} more '
-                        f'time{"s" if count > 1 else ""}]\n'
-                    )
-                last_file = frame.filename
-                last_line = frame.lineno
-                last_name = frame.name
-                count = 0
-            count += 1
-            if count > _RECURSIVE_CUTOFF:
-                continue
-            result.append(self.format_frame(frame))
+                    continue
+                result.append(formatted_frame)
 
         if count > _RECURSIVE_CUTOFF:
             count -= _RECURSIVE_CUTOFF
@@ -618,7 +621,7 @@ class TracebackException:
 
     def __init__(self, exc_type, exc_value, exc_traceback, *, limit=None,
             lookup_lines=True, capture_locals=False, compact=False,
-            _seen=None):
+            stack_summary_cls=None, _seen=None):
         # NB: we need to accept exc_traceback, exc_value, exc_traceback to
         # permit backwards compat with the existing API, otherwise we
         # need stub thunk objects just to glue it together.
@@ -628,8 +631,9 @@ def __init__(self, exc_type, exc_value, exc_traceback, *, limit=None,
             _seen = set()
         _seen.add(id(exc_value))
 
-        # TODO: locals.
-        self.stack = StackSummary._extract_from_extended_frame_gen(
+        if stack_summary_cls is None:
+            stack_summary_cls = StackSummary
+        self.stack = stack_summary_cls._extract_from_extended_frame_gen(
             _walk_tb_with_full_positions(exc_traceback),
             limit=limit, lookup_lines=lookup_lines,
             capture_locals=capture_locals)
@@ -665,6 +669,7 @@ def __init__(self, exc_type, exc_value, exc_traceback, *, limit=None,
                         limit=limit,
                         lookup_lines=lookup_lines,
                         capture_locals=capture_locals,
+                        stack_summary_cls=stack_summary_cls,
                         _seen=_seen)
                 else:
                     cause = None
@@ -684,6 +689,7 @@ def __init__(self, exc_type, exc_value, exc_traceback, *, limit=None,
                         limit=limit,
                         lookup_lines=lookup_lines,
                         capture_locals=capture_locals,
+                        stack_summary_cls=stack_summary_cls,
                         _seen=_seen)
                 else:
                     context = None

Misc/NEWS.d/next/Library/2021-06-17-16-16-33.bpo-31299.d4WDz7.rst

@@ -0,0 +1 @@
+Added the ``stack_summary_cls`` parameter to :class:`TracebackException`, to allow fine-grained control over the content of a formatted traceback.  Added option to completely drop frames from the output by returning ``None`` from a :meth:`~StackSummary.format_frame` override.