Post-rebase fixes

line_profiler/_line_profiler.pyx
    _SysMonitoringState.line_tracing_events, .line_tracing_event_set
        New class attributes
    _LineProfilerManager
        _base_callback()
            - Updated call signature
            - Now using `._call_callback()` to handle stored callbacks
        _call_callback()
            New helper method for calling the stored callback which:
            - Fetches changes to the `sys.monitoring` events and
              callables to `self.mon_state`
            - Restores the global event list and tool-id lock to ensure
              continued profiling if necessary

tests/test_sys_trace.py
    suspend_tracing
        Updated implementation to also suspend `sys.monitoring`-based
        line profiling
    _test_helper_callback_preservation()
        Updated to skip the legacy-tracing check when using
        `sys.monitoring`-based line profiling
    test_callback_wrapping()
        Updated implementation to reflect non-interference with legacy
        trace functions when using `sys.monitoring`-based line profiling
    test_wrapping_thread_local_callbacks()
        - Updated parametrization to reflect non-interference with
          legacy trace functions when using `sys.monitoring`-based line
          profiling
        - Updated assertion to give more informative errors

Author: TTsangSC

line_profiler/_line_profiler.pyx

@@ -314,6 +314,18 @@ cdef class _SysMonitoringState:
     cdef dict callbacks  # type: dict[int, Callable | None]
     cdef int events
 
+    if CAN_USE_SYS_MONITORING:
+        line_tracing_event_set = (  # type: ClassVar[FrozenSet[int]]
+            frozenset({sys.monitoring.events.LINE,
+                       sys.monitoring.events.PY_RETURN,
+                       sys.monitoring.events.PY_YIELD}))
+        line_tracing_events = (sys.monitoring.events.LINE
+                               | sys.monitoring.events.PY_RETURN
+                               | sys.monitoring.events.PY_YIELD)
+    else:
+        line_tracing_event_set = frozenset({})
+        line_tracing_events = 0
+
     def __init__(self, name=None, callbacks=None, events=0):
         self.name = name
         self.callbacks = callbacks or {}
@@ -342,11 +354,7 @@ cdef class _SysMonitoringState:
             self.events = mon.get_events(mon.PROFILER_ID)
             mon.free_tool_id(mon.PROFILER_ID)
         mon.use_tool_id(mon.PROFILER_ID, 'line_profiler')
-        events = (self.events
-                  | mon.events.LINE
-                  | mon.events.PY_RETURN
-                  | mon.events.PY_YIELD)
-        mon.set_events(mon.PROFILER_ID, events)
+        mon.set_events(mon.PROFILER_ID, self.events | self.line_tracing_events)
 
         # Register tracebacks
         for event_id, callback in [
@@ -435,7 +443,7 @@ cdef class _LineProfilerManager:
         :py:func:`sys.monitoring.register_callback`.
         """
         self._base_callback(
-            sys.monitoring.events.LINE, code, lineno, (lineno,))
+            1, sys.monitoring.events.LINE, code, lineno, (lineno,))
 
     @cython.profile(False)
     cpdef handle_return_event(
@@ -467,47 +475,69 @@ cdef class _LineProfilerManager:
             This is deliberately made a non-traceable C method so that
             we don't fall info infinite recursion.
         """
-        self._base_callback(event_id,
-                            code,
-                            PyCode_Addr2Line(<PyCodeObject*>code, offset),
-                            (offset, retval))
+        cdef int lineno = PyCode_Addr2Line(<PyCodeObject*>code, offset)
+        self._base_callback(0, event_id, code, lineno, (offset, retval))
 
     cdef void _base_callback(
-            self, int event_id, object code, int lineno, object other_args):
+            self, int is_line_event, int event_id,
+            object code, int lineno, object other_args):
         """
         Base for the various callbacks passed to
-        :py:func:`sys.monitoring.register_callback`.  Also takes care of
-        the restoration of callbacks should they be unset.
+        :py:func:`sys.monitoring.register_callback`.
 
         Note:
             This is deliberately made a non-traceable C method so that
             we don't fall info infinite recursion.
         """
-        cdef object callback_before, callback_after, callback_wrapped
-        cdef dict callbacks = self.mon_state.callbacks
-        mon = sys.monitoring
-        inner_trace_callback((event_id == mon.events.LINE),
-                             self.active_instances,
-                             code,
-                             lineno)
-        if not self.wrap_trace:
-            return
+        inner_trace_callback(
+            is_line_event, self.active_instances, code, lineno)
+        if self._wrap_trace:
+            self._call_callback(event_id, code, other_args)
+
+    cdef void _call_callback(
+            self, int event_id, object code, object other_args):
+        """
+        Call the stored callback in ``self.mon_state``.  Also takes care
+        of the restoration of :py:mod:`sys.monitoring` callbacks,
+        tool-id lock, and events should they be unset.
 
-        # Call wrapped callback
-        callback_wrapped = callbacks.get(event_id)
-        if callback_wrapped is None:
+        Note:
+            This is deliberately made a non-traceable C method so that
+            we don't fall info infinite recursion.
+        """
+        mon = sys.monitoring
+        cdef object callback  # type: Callable | None
+        cdef object callback_after  # type: Callable | None
+        cdef object callback_wrapped  # type: Callable | None
+        cdef int events_needed
+        cdef int ev_id
+        cdef int prof_id = mon.PROFILER_ID
+        cdef _SysMonitoringState state = self.mon_state
+        cdef dict callbacks_before = {}
+        # Call wrapped callback where suitable
+        callback_wrapped = state.callbacks.get(event_id)
+        if callback_wrapped is None or not (event_id & state.events):
             return
-        callback_before = get_current_callback(event_id)
+        for ev_id in state.line_tracing_event_set:
+            callbacks_before[ev_id] = get_current_callback(ev_id)
+        events_needed = state.line_tracing_events
         try:
             callback_wrapped(code, *other_args)
         finally:
-            callback_after = get_current_callback(event_id)
-            if callback_after is None:
-                # The wrapped callback has unset itself;
-                # remove from `.mon_state.callbacks`
-                callbacks[event_id] = None
-            if callback_before is not callback_after:
-                mon.register_callback(mon.PROFILER_ID, callback_before)
+            state.events = mon.get_events(prof_id)
+            register = mon.register_callback
+            # If the wrapped callback has changed:
+            for ev_id, callback in callbacks_before.items():
+                # - Restore the `sys.monitoring` callback
+                callback_after = register(prof_id, ev_id, callback)
+                # - Remember the updated callback in `state.callbacks`
+                if callback is not callback_after:
+                    state.callbacks[ev_id] = callback_after
+            # Reset the tool ID lock if released
+            if not mon.get_tool(prof_id):
+                mon.use_tool_id(prof_id, 'line_profiler')
+            # Restore the `sys.monitoring` events if unset
+            mon.set_events(prof_id, state.events | events_needed)
 
     cpdef _handle_enable_event(self, prof):
         cdef TraceCallback* legacy_callback

tests/test_sys_trace.py

@@ -26,14 +26,15 @@
 from ast import literal_eval
 from contextlib import nullcontext
 from io import StringIO
-from types import FrameType
+from types import FrameType, ModuleType
 from typing import Any, Optional, Union, Callable, List, Literal
 from line_profiler import LineProfiler
 
 
 # Common utilities
 
 DEBUG = False
+USE_SYS_MONITORING = isinstance(getattr(sys, 'monitoring', None), ModuleType)
 
 Event = Literal['call', 'line', 'return', 'exception', 'opcode']
 TracingFunc = Callable[[FrameType, Event, Any], Union['TracingFunc', None]]
@@ -163,14 +164,25 @@ def baz(n: int) -> int:
 class suspend_tracing:
     def __init__(self):
         self.callback = None
+        self.events = 0
 
     def __enter__(self):
-        self.callback = sys.gettrace()
-        sys.settrace(None)
+        if USE_SYS_MONITORING:
+            mod = sys.monitoring
+            self.events = mod.get_events(mod.PROFILER_ID)
+            mod.set_events(mod.PROFILER_ID, mod.events.NO_EVENTS)
+        else:
+            self.callback = sys.gettrace()
+            sys.settrace(None)
 
     def __exit__(self, *_, **__):
-        sys.settrace(self.callback)
-        self.callback = None
+        if USE_SYS_MONITORING:
+            mod = sys.monitoring
+            mod.set_events(mod.PROFILER_ID, self.events)
+            self.events = 0
+        else:
+            sys.settrace(self.callback)
+            self.callback = None
 
 
 def get_incr_logger(logs: List[str], func: Literal[foo, bar, baz] = foo, *,
@@ -270,8 +282,9 @@ def _test_helper_callback_preservation(
     assert sys.gettrace() is callback, f'can\'t set trace to {callback!r}'
     profile = LineProfiler(wrap_trace=False)
     profile.enable_by_count()
-    assert profile in sys.gettrace().active_instances, (
-        'can\'t set trace to the profiler')
+    if not USE_SYS_MONITORING:
+        assert profile in sys.gettrace().active_instances, (
+            'can\'t set trace to the profiler')
     profile.disable_by_count()
     assert sys.gettrace() is callback, f'trace not restored to {callback!r}'
     sys.settrace(None)
@@ -311,7 +324,7 @@ def test_callback_wrapping(
     else:
         foo_like = foo
         trace_preserved = True
-    if trace_preserved:
+    if trace_preserved or USE_SYS_MONITORING:
         exp_logs = [f'foo: spam = {spam}' for spam in range(1, 6)]
     else:
         exp_logs = []
@@ -558,9 +571,12 @@ def test_wrapping_thread_local_callbacks(label: str,
     [(True, True, 100, {0: 2,  # Both calls are traced
                         5: 0,  # Tracing suspended
                         7: 100}),  # Tracing restored (both calls)
-     # If `set_frame_local_trace` is false, tracing is suspended for the
-     # rest of the frame
-     (True, False, 100, {0: 2, 5: 0, 7: 0}),
+     # If `set_frame_local_trace` is false:
+     # - When using legacy tracing, tracing is suspended for the rest of
+     #   the frame
+     # - Else, tracing is unaffected
+     (True, False, 100,
+      {0: 2, 5: 0, 7: 100 if USE_SYS_MONITORING else 0}),
      # Calling a function always triggers `<trace>.__call__()`
      (False, True, 100, {0: 1,  # Only one of the calls is traced
                          2: 100}),  # 100 hits on the line in the loop
@@ -621,4 +637,4 @@ def func_break_in_middle(n):
     all_nhits = {lineno - body_start_line: _nhits
                  for (lineno, _nhits, _) in entries}
     all_nhits = {lineno: all_nhits.get(lineno, 0) for lineno in nhits}
-    assert all_nhits == nhits
+    assert all_nhits == nhits, f'expected {nhits=}, got {all_nhits=}'