Add python impl of timing thread

pyinstrument/low_level/pyi_timing_thread_python.py

@@ -73,7 +73,7 @@ def pyi_timing_thread_unsubscribe(id: int):
                 thread_alive = False
             return 0
         else:
-            return "PYI_TIMING_THREAD_NOT_SUBSCRIBED"
+            raise Exception("PYI_TIMING_THREAD_NOT_SUBSCRIBED")
 
 
 def pyi_timing_thread_get_time() -> float:

pyinstrument/low_level/stat_profile_python.py

@@ -6,34 +6,67 @@
 import types
 from typing import Any, Callable, List, Optional, Type
 
+from pyinstrument.low_level.pyi_timing_thread_python import (
+    pyi_timing_thread_get_time,
+    pyi_timing_thread_subscribe,
+    pyi_timing_thread_unsubscribe,
+)
+from pyinstrument.low_level.types import TimerType
+
 
 class PythonStatProfiler:
     await_stack: list[str]
-
-    def __init__(self, target, interval, context_var, timer_func):
+    timing_thread_subscription: int | None = None
+
+    def __init__(
+        self,
+        target: Callable[[types.FrameType, str, Any], Any],
+        interval: float,
+        context_var: contextvars.ContextVar[object | None] | None,
+        timer_type: TimerType,
+        timer_func: Callable[[], float] | None,
+    ):
         self.target = target
         self.interval = interval
-        self.timer_func = timer_func or timeit.default_timer
-        self.last_invocation = self.timer_func()
-
         if context_var:
             # raise typeerror to match the C version
             if not isinstance(context_var, contextvars.ContextVar):
                 raise TypeError("not a context var")
-
         self.context_var = context_var
+
+        self.timer_type = timer_type
+
+        if timer_type == "walltime":
+            self.get_time = timeit.default_timer
+        elif timer_type == "walltime_thread":
+            self.get_time = pyi_timing_thread_get_time
+            self.timing_thread_subscription = pyi_timing_thread_subscribe(interval)
+        elif timer_type == "timer_func":
+            if timer_func is None:
+                raise TypeError("timer_func must be provided for timer_func timer_type")
+            self.get_time = timer_func
+        else:
+            raise ValueError(f"invalid timer_type '{timer_type}'")
+
+        self.last_invocation = self.get_time()
+
         self.last_context_var_value = context_var.get() if context_var else None
         self.await_stack = []
 
+    def __del__(self):
+        if self.timing_thread_subscription is not None:
+            pyi_timing_thread_unsubscribe(self.timing_thread_subscription)
+
     def profile(self, frame: types.FrameType, event: str, arg: Any):
-        now = self.timer_func()
+        now = self.get_time()
 
         if self.context_var:
             context_var_value = self.context_var.get()
             last_context_var_value = self.last_context_var_value
 
             if context_var_value is not last_context_var_value:
                 context_change_frame = frame.f_back if event == "call" else frame
+                assert context_change_frame is not None
                 self.target(
                     context_change_frame,
                     "context_changed",
@@ -64,13 +97,15 @@ def setstatprofile(
     target: Callable[[types.FrameType, str, Any], Any] | None,
     interval: float = 0.001,
     context_var: contextvars.ContextVar[object | None] | None = None,
+    timer_type: TimerType = "walltime",
     timer_func: Callable[[], float] | None = None,
 ) -> None:
     if target:
         profiler = PythonStatProfiler(
             target=target,
             interval=interval,
             context_var=context_var,
+            timer_type=timer_type,
             timer_func=timer_func,
         )
         sys.setprofile(profiler.profile)