Support a new --native-last argument (bloomberg#182)

Add a new ``--native-last`` command line flag to only show the C frames
after the last Python frame on each thread.

Author: Kretzer6019

news/182.feature.rst

@@ -0,0 +1 @@
+Add a new ``--native-last`` command line flag to only show the C frames after the last Python frame.

src/pystack/__main__.py

@@ -166,6 +166,15 @@ def generate_cli_parser() -> argparse.ArgumentParser:
         help="Include native (C) frames from threads not registered with "
         "the interpreter (implies --native)",
     )
+    remote_parser.add_argument(
+        "--native-last",
+        action="store_const",
+        dest="native_mode",
+        const=NativeReportingMode.LAST,
+        default=NativeReportingMode.OFF,
+        help="Include native (C) frames only after the last python frame "
+        "in the resulting stack trace",
+    )
     remote_parser.add_argument(
         "--locals",
         action="store_true",
@@ -212,6 +221,15 @@ def generate_cli_parser() -> argparse.ArgumentParser:
         help="Include native (C) frames from threads not registered with "
         "the interpreter (implies --native)",
     )
+    core_parser.add_argument(
+        "--native-last",
+        action="store_const",
+        dest="native_mode",
+        const=NativeReportingMode.LAST,
+        default=NativeReportingMode.OFF,
+        help="Include native (C) frames only after the last python frame "
+        "in the resulting stack trace",
+    )
     core_parser.add_argument(
         "--locals",
         action="store_true",
@@ -274,8 +292,7 @@ def process_remote(parser: argparse.ArgumentParser, args: argparse.Namespace) ->
         locals=args.locals,
         method=StackMethod.ALL if args.exhaustive else StackMethod.AUTO,
     ):
-        native = args.native_mode != NativeReportingMode.OFF
-        print_thread(thread, native)
+        print_thread(thread, args.native_mode)
 
 
 def format_psinfo_information(psinfo: Dict[str, Any]) -> str:
@@ -405,8 +422,7 @@ def process_core(parser: argparse.ArgumentParser, args: argparse.Namespace) -> N
         locals=args.locals,
         method=StackMethod.ALL if args.exhaustive else StackMethod.AUTO,
     ):
-        native = args.native_mode != NativeReportingMode.OFF
-        print_thread(thread, native)
+        print_thread(thread, args.native_mode)
 
 
 if __name__ == "__main__":  # pragma: no cover

src/pystack/_pystack.pyi

@@ -24,9 +24,10 @@ class CoreFileAnalyzer:
     def missing_modules(self) -> List[str]: ...
 
 class NativeReportingMode(enum.Enum):
-    ALL = 1
-    OFF = 2
-    PYTHON = 3
+    ALL = ...
+    OFF = ...
+    PYTHON = ...
+    LAST = ...
 
 class StackMethod(enum.Enum):
     ALL = 1

src/pystack/_pystack.pyx

@@ -86,6 +86,7 @@ class NativeReportingMode(enum.Enum):
     OFF = 0
     PYTHON = 1
     ALL = 1000
+    LAST = 2000
 
 
 cdef api void log_with_python(const cppstring *message, int level) noexcept:

src/pystack/traceback_formatter.py

@@ -4,15 +4,16 @@
 from typing import Optional
 
 from .colors import colored
+from .engine import NativeReportingMode
 from .types import NativeFrame
 from .types import PyCodeObject
 from .types import PyFrame
 from .types import PyThread
 from .types import frame_type
 
 
-def print_thread(thread: PyThread, native: bool) -> None:
-    for line in format_thread(thread, native):
+def print_thread(thread: PyThread, native_mode: NativeReportingMode) -> None:
+    for line in format_thread(thread, native_mode):
         print(line, file=sys.stdout, flush=True)
 
 
@@ -62,7 +63,8 @@ def _are_the_stacks_mergeable(thread: PyThread) -> bool:
     return n_eval_frames == n_entry_frames
 
 
-def format_thread(thread: PyThread, native: bool) -> Iterable[str]:
+def format_thread(thread: PyThread, native_mode: NativeReportingMode) -> Iterable[str]:
+    native = native_mode != NativeReportingMode.OFF
     current_frame: Optional[PyFrame] = thread.first_frame
     if current_frame is None and not native:
         yield f"The frame stack for thread {thread.tid} is empty"
@@ -82,16 +84,22 @@ def format_thread(thread: PyThread, native: bool) -> Iterable[str]:
                 yield from format_frame(current_frame)
             current_frame = current_frame.next
     else:
-        yield from _format_merged_stacks(thread, current_frame)
+        yield from _format_merged_stacks(
+            thread, current_frame, native_mode == NativeReportingMode.LAST
+        )
     yield ""
 
 
 def _format_merged_stacks(
-    thread: PyThread, current_frame: Optional[PyFrame]
+    thread: PyThread,
+    current_frame: Optional[PyFrame],
+    native_last: bool = False,
 ) -> Iterable[str]:
+    c_frames_list: list[str] = []
     for frame in thread.native_frames:
         if frame_type(frame, thread.python_version) == NativeFrame.FrameType.EVAL:
             assert current_frame is not None
+            c_frames_list = []
             yield from format_frame(current_frame)
             current_frame = current_frame.next
             while current_frame and not current_frame.is_entry:
@@ -102,12 +110,18 @@ def _format_merged_stacks(
             continue
         elif frame_type(frame, thread.python_version) == NativeFrame.FrameType.OTHER:
             function = colored(frame.symbol, "yellow")
-            yield (
+            formatted_c_frame = (
                 f'    {colored("(C)", "blue")} File "{frame.path}",'
                 f" line {frame.linenumber},"
                 f" in {function} ({colored(frame.library, attrs=['faint'])})"
             )
+            if native_last:
+                c_frames_list.append(formatted_c_frame)
+            else:
+                yield formatted_c_frame
         else:  # pragma: no cover
             raise ValueError(
                 f"Invalid frame type: {frame_type(frame, thread.python_version)}"
             )
+    for c_frame in c_frames_list:
+        yield c_frame

tests/unit/test_main.py

@@ -207,7 +207,9 @@ def test_process_remote_default():
         locals=False,
         method=StackMethod.AUTO,
     )
-    assert print_thread_mock.mock_calls == [call(thread, False) for thread in threads]
+    assert print_thread_mock.mock_calls == [
+        call(thread, NativeReportingMode.OFF) for thread in threads
+    ]
 
 
 def test_process_remote_no_block():
@@ -238,13 +240,23 @@ def test_process_remote_no_block():
         locals=False,
         method=StackMethod.AUTO,
     )
-    assert print_thread_mock.mock_calls == [call(thread, False) for thread in threads]
+    assert print_thread_mock.mock_calls == [
+        call(thread, NativeReportingMode.OFF) for thread in threads
+    ]
 
 
-def test_process_remote_native():
+@pytest.mark.parametrize(
+    "argument, mode",
+    [
+        ["--native", NativeReportingMode.PYTHON],
+        ["--native-all", NativeReportingMode.ALL],
+        ["--native-last", NativeReportingMode.LAST],
+    ],
+)
+def test_process_remote_native(argument, mode):
     # GIVEN
 
-    argv = ["pystack", "remote", "31", "--native"]
+    argv = ["pystack", "remote", "31", argument]
 
     threads = [Mock(), Mock(), Mock()]
 
@@ -265,11 +277,11 @@ def test_process_remote_native():
     get_process_threads_mock.assert_called_with(
         31,
         stop_process=True,
-        native_mode=NativeReportingMode.PYTHON,
+        native_mode=mode,
         locals=False,
         method=StackMethod.AUTO,
     )
-    assert print_thread_mock.mock_calls == [call(thread, True) for thread in threads]
+    assert print_thread_mock.mock_calls == [call(thread, mode) for thread in threads]
 
 
 def test_process_remote_locals():
@@ -300,7 +312,9 @@ def test_process_remote_locals():
         locals=True,
         method=StackMethod.AUTO,
     )
-    assert print_thread_mock.mock_calls == [call(thread, False) for thread in threads]
+    assert print_thread_mock.mock_calls == [
+        call(thread, NativeReportingMode.OFF) for thread in threads
+    ]
 
 
 def test_process_remote_native_no_block(capsys):
@@ -357,7 +371,9 @@ def test_process_remote_exhaustive():
         locals=False,
         method=StackMethod.ALL,
     )
-    assert print_thread_mock.mock_calls == [call(thread, False) for thread in threads]
+    assert print_thread_mock.mock_calls == [
+        call(thread, NativeReportingMode.OFF) for thread in threads
+    ]
 
 
 @pytest.mark.parametrize(
@@ -432,7 +448,9 @@ def test_process_core_default_without_executable():
         locals=False,
         method=StackMethod.AUTO,
     )
-    assert print_thread_mock.mock_calls == [call(thread, False) for thread in threads]
+    assert print_thread_mock.mock_calls == [
+        call(thread, NativeReportingMode.OFF) for thread in threads
+    ]
 
 
 def test_process_core_default_gzip_without_executable():
@@ -486,7 +504,9 @@ def test_process_core_default_gzip_without_executable():
         locals=False,
         method=StackMethod.AUTO,
     )
-    assert print_thread_mock.mock_calls == [call(thread, False) for thread in threads]
+    assert print_thread_mock.mock_calls == [
+        call(thread, NativeReportingMode.OFF) for thread in threads
+    ]
     gzip_open_mock.assert_called_with(Path("corefile.gz"), "rb")
 
 
@@ -580,14 +600,17 @@ def test_process_core_default_with_executable():
         locals=False,
         method=StackMethod.AUTO,
     )
-    assert print_thread_mock.mock_calls == [call(thread, False) for thread in threads]
+    assert print_thread_mock.mock_calls == [
+        call(thread, NativeReportingMode.OFF) for thread in threads
+    ]
 
 
 @pytest.mark.parametrize(
     "argument, mode",
     [
         ["--native", NativeReportingMode.PYTHON],
         ["--native-all", NativeReportingMode.ALL],
+        ["--native-last", NativeReportingMode.LAST],
     ],
 )
 def test_process_core_native(argument, mode):
@@ -627,7 +650,7 @@ def test_process_core_native(argument, mode):
         locals=False,
         method=StackMethod.AUTO,
     )
-    assert print_thread_mock.mock_calls == [call(thread, True) for thread in threads]
+    assert print_thread_mock.mock_calls == [call(thread, mode) for thread in threads]
 
 
 def test_process_core_locals():
@@ -667,7 +690,9 @@ def test_process_core_locals():
         locals=True,
         method=StackMethod.AUTO,
     )
-    assert print_thread_mock.mock_calls == [call(thread, False) for thread in threads]
+    assert print_thread_mock.mock_calls == [
+        call(thread, NativeReportingMode.OFF) for thread in threads
+    ]
 
 
 def test_process_core_with_search_path():
@@ -714,7 +739,9 @@ def test_process_core_with_search_path():
         locals=False,
         method=StackMethod.AUTO,
     )
-    assert print_thread_mock.mock_calls == [call(thread, False) for thread in threads]
+    assert print_thread_mock.mock_calls == [
+        call(thread, NativeReportingMode.OFF) for thread in threads
+    ]
 
 
 def test_process_core_with_search_root():
@@ -762,7 +789,9 @@ def test_process_core_with_search_root():
         locals=False,
         method=StackMethod.AUTO,
     )
-    assert print_thread_mock.mock_calls == [call(thread, False) for thread in threads]
+    assert print_thread_mock.mock_calls == [
+        call(thread, NativeReportingMode.OFF) for thread in threads
+    ]
 
 
 def test_process_core_with_not_readable_search_root():
@@ -947,7 +976,9 @@ def test_process_core_exhaustive():
         locals=False,
         method=StackMethod.ALL,
     )
-    assert print_thread_mock.mock_calls == [call(thread, False) for thread in threads]
+    assert print_thread_mock.mock_calls == [
+        call(thread, NativeReportingMode.OFF) for thread in threads
+    ]
 
 
 def test_default_colored_output():