Merge pull request #347 from TTsangSC/multi-profiler

ENH: use multiple profiler instances

Author: Erotemic

CHANGELOG.rst

@@ -14,6 +14,7 @@ Changes
 * ENH: Added CLI arguments ``-c`` to ``kernprof`` for (auto-)profiling module/package/inline-script execution instead of that of script files; passing ``'-'`` as the script-file name now also reads from and profiles ``stdin``
 * ENH: In Python >=3.11, profiled objects are reported using their qualified name.
 * ENH: Highlight final summary using rich if enabled
+* ENH: Made it possible to use multiple profiler instances simultaneously
 
 4.2.0
 ~~~~~

line_profiler/_line_profiler.pyx

@@ -253,6 +253,9 @@ cdef class LineProfiler:
     cdef public double timer_unit
     cdef public object threaddata
 
+    # This is shared between instances and threads
+    _all_active_instances = {}
+
     def __init__(self, *functions):
         self.functions = []
         self.code_hash_map = {}
@@ -319,6 +322,16 @@ cdef class LineProfiler:
         def __set__(self, value):
             self.threaddata.enable_count = value
 
+    # This is shared between instances, but thread-local
+    property _active_instances:
+        def __get__(self):
+            thread_id = threading.get_ident()
+            try:
+                return self._all_active_instances[thread_id]
+            except KeyError:
+                insts = self._all_active_instances[thread_id] = set()
+                return insts
+
     def enable_by_count(self):
         """ Enable the profiler if it hasn't been enabled before.
         """
@@ -345,8 +358,11 @@ cdef class LineProfiler:
         # Register `line_profiler` with `sys.monitoring` in Python 3.12
         # and above;
         # see: https://docs.python.org/3/library/sys.monitoring.html
-        _sys_monitoring_register()
-        PyEval_SetTrace(python_trace_callback, self)
+        instances = self._active_instances
+        if not instances:
+            _sys_monitoring_register()
+            PyEval_SetTrace(python_trace_callback, instances)
+        instances.add(self)
 
     @property
     def c_code_map(self):
@@ -397,12 +413,15 @@ cdef class LineProfiler:
 
 
     cpdef disable(self):
+        instances = self._active_instances
         self._c_last_time[threading.get_ident()].clear()
-        unset_trace()
-        # Deregister `line_profiler` with `sys.monitoring` in Python
-        # 3.12 and above;
-        # see: https://docs.python.org/3/library/sys.monitoring.html
-        _sys_monitoring_deregister()
+        instances.discard(self)
+        if not instances:
+            unset_trace()
+            # Deregister `line_profiler` with `sys.monitoring` in Python
+            # 3.12 and above;
+            # see: https://docs.python.org/3/library/sys.monitoring.html
+            _sys_monitoring_deregister()
 
     def get_stats(self):
         """
@@ -443,53 +462,59 @@ cdef class LineProfiler:
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cdef extern int python_trace_callback(object self_, PyFrameObject *py_frame,
+cdef extern int python_trace_callback(object instances,
+                                      PyFrameObject *py_frame,
                                       int what, PyObject *arg):
     """
     The PyEval_SetTrace() callback.
 
     References:
        https://github.com/python/cpython/blob/de2a4036/Include/cpython/pystate.h#L16 
     """
-    cdef LineProfiler self
+    cdef object prof_
+    cdef LineProfiler prof
     cdef object code
     cdef LineTime entry
     cdef LastTime old
     cdef int key
     cdef PY_LONG_LONG time
+    cdef int has_time = 0
     cdef int64 code_hash
     cdef int64 block_hash
     cdef unordered_map[int64, LineTime] line_entries
     cdef uint64 linenum
 
-    self = <LineProfiler>self_
-
     if what == PyTrace_LINE or what == PyTrace_RETURN:
         # Normally we'd need to DECREF the return from get_frame_code, but Cython does that for us
         block_hash = hash(get_frame_code(py_frame))
 
         linenum = PyFrame_GetLineNumber(py_frame)
         code_hash = compute_line_hash(block_hash, linenum)
 
-        if self._c_code_map.count(code_hash):
-            time = hpTimer()
+        for prof_ in instances:
+            prof = <LineProfiler>prof_
+            if not prof._c_code_map.count(code_hash):
+                continue
+            if not has_time:
+                time = hpTimer()
+                has_time = 1
             ident = threading.get_ident()
-            if self._c_last_time[ident].count(block_hash):
-                old = self._c_last_time[ident][block_hash]
-                line_entries = self._c_code_map[code_hash]
+            if prof._c_last_time[ident].count(block_hash):
+                old = prof._c_last_time[ident][block_hash]
+                line_entries = prof._c_code_map[code_hash]
                 key = old.f_lineno
                 if not line_entries.count(key):
-                    self._c_code_map[code_hash][key] = LineTime(code_hash, key, 0, 0)
-                self._c_code_map[code_hash][key].nhits += 1
-                self._c_code_map[code_hash][key].total_time += time - old.time
+                    prof._c_code_map[code_hash][key] = LineTime(code_hash, key, 0, 0)
+                prof._c_code_map[code_hash][key].nhits += 1
+                prof._c_code_map[code_hash][key].total_time += time - old.time
             if what == PyTrace_LINE:
                 # Get the time again. This way, we don't record much time wasted
                 # in this function.
-                self._c_last_time[ident][block_hash] = LastTime(linenum, hpTimer())
-            elif self._c_last_time[ident].count(block_hash):
+                prof._c_last_time[ident][block_hash] = LastTime(linenum, hpTimer())
+            elif prof._c_last_time[ident].count(block_hash):
                 # We are returning from a function, not executing a line. Delete
                 # the last_time record. It may have already been deleted if we
                 # are profiling a generator that is being pumped past its end.
-                self._c_last_time[ident].erase(self._c_last_time[ident].find(block_hash))
+                prof._c_last_time[ident].erase(prof._c_last_time[ident].find(block_hash))
 
     return 0

line_profiler/line_profiler.py

@@ -4,12 +4,12 @@
 inspect its output. This depends on the :py:mod:`line_profiler._line_profiler`
 Cython backend.
 """
-import pickle
 import inspect
 import linecache
-import tempfile
 import os
+import pickle
 import sys
+import tempfile
 from argparse import ArgumentError, ArgumentParser
 
 try:
@@ -66,6 +66,21 @@ def _get_underlying_functions(func):
     return [type(func).__call__]
 
 
+class _WrapperInfo:
+    """
+    Helper object for holding the state of a wrapper function.
+
+    Attributes:
+        func (types.FunctionType):
+            The function it wraps.
+        profiler_id (int)
+            ID of the `LineProfiler`.
+    """
+    def __init__(self, func, profiler_id):
+        self.func = func
+        self.profiler_id = profiler_id
+
+
 class LineProfiler(CLineProfiler, ByCountProfilerMixin):
     """
     A profiler that records the execution times of individual lines.
@@ -83,7 +98,6 @@ class LineProfiler(CLineProfiler, ByCountProfilerMixin):
         >>> func()
         >>> profile.print_stats()
     """
-
     def __call__(self, func):
         """
         Decorate a function, method, property, partial object etc. to
@@ -104,12 +118,15 @@ def add_callable(self, func):
         Returns:
             1 if any function is added to the profiler, 0 otherwise
         """
-        guard = self._already_wrapped
-
         nadded = 0
         for impl in _get_underlying_functions(func):
-            if guard(impl):
+            info, wrapped_by_this_prof = self._get_wrapper_info(impl)
+            if wrapped_by_this_prof:
                 continue
+            if info:
+                # It's still a profiling wrapper, just wrapped by
+                # someone else -> extract the inner function
+                impl = info.func
             self.add_function(impl)
             nadded += 1
 
@@ -150,6 +167,25 @@ def add_module(self, mod):
 
         return nfuncsadded
 
+    def _get_wrapper_info(self, func):
+        info = getattr(func, self._profiler_wrapped_marker, None)
+        return info, bool(info and id(self) == info.profiler_id)
+
+    # Override these mixed-in bookkeeping methods to take care of
+    # potential multiple profiler sequences
+
+    def _already_a_wrapper(self, func):
+        return self._get_wrapper_info(func)[1]
+
+    def _mark_wrapper(self, wrapper):
+        # Are re-wrapping an existing wrapper (e.g. created by another
+        # profiler?)
+        wrapped = wrapper.__wrapped__
+        info = getattr(wrapped, self._profiler_wrapped_marker, None)
+        new_info = _WrapperInfo(info.func if info else wrapped, id(self))
+        setattr(wrapper, self._profiler_wrapped_marker, new_info)
+        return wrapper
+
 
 # This could be in the ipython_extension submodule,
 # but it doesn't depend on the IPython module so it's easier to just let it stay here.

line_profiler/line_profiler.pyi

@@ -1,5 +1,4 @@
-from typing import List
-from typing import Tuple
+from typing import Literal, List, Tuple
 import io
 from ._line_profiler import LineProfiler as CLineProfiler
 from .profiler_mixin import ByCountProfilerMixin
@@ -11,8 +10,7 @@ def load_ipython_extension(ip) -> None:
 
 
 class LineProfiler(CLineProfiler, ByCountProfilerMixin):
-
-    def add_callable(self, func) -> None:
+    def add_callable(self, func) -> Literal[0, 1]:
         ...
 
     def dump_stats(self, filename) -> None:
@@ -28,7 +26,7 @@ class LineProfiler(CLineProfiler, ByCountProfilerMixin):
                     rich: bool = ...) -> None:
         ...
 
-    def add_module(self, mod):
+    def add_module(self, mod) -> int:
         ...
 
 

line_profiler/profiler_mixin.py

@@ -198,7 +198,7 @@ def wrap_async_generator(self, func):
         Wrap an async generator function to profile it.
         """
         # Prevent double-wrap
-        if self._already_wrapped(func):
+        if self._already_a_wrapper(func):
             return func
 
         @functools.wraps(func)
@@ -216,14 +216,14 @@ async def wrapper(*args, **kwds):
                     self.disable_by_count()
                 input_ = (yield item)
 
-        return self._mark_wrapped(wrapper)
+        return self._mark_wrapper(wrapper)
 
     def wrap_coroutine(self, func):
         """
         Wrap a coroutine function to profile it.
         """
         # Prevent double-wrap
-        if self._already_wrapped(func):
+        if self._already_a_wrapper(func):
             return func
 
         @functools.wraps(func)
@@ -235,14 +235,14 @@ async def wrapper(*args, **kwds):
                 self.disable_by_count()
             return result
 
-        return self._mark_wrapped(wrapper)
+        return self._mark_wrapper(wrapper)
 
     def wrap_generator(self, func):
         """
         Wrap a generator function to profile it.
         """
         # Prevent double-wrap
-        if self._already_wrapped(func):
+        if self._already_a_wrapper(func):
             return func
 
         @functools.wraps(func)
@@ -260,14 +260,14 @@ def wrapper(*args, **kwds):
                     self.disable_by_count()
                 input_ = (yield item)
 
-        return self._mark_wrapped(wrapper)
+        return self._mark_wrapper(wrapper)
 
     def wrap_function(self, func):
         """
         Wrap a function to profile it.
         """
         # Prevent double-wrap
-        if self._already_wrapped(func):
+        if self._already_a_wrapper(func):
             return func
 
         @functools.wraps(func)
@@ -279,14 +279,14 @@ def wrapper(*args, **kwds):
                 self.disable_by_count()
             return result
 
-        return self._mark_wrapped(wrapper)
+        return self._mark_wrapper(wrapper)
 
-    def _already_wrapped(self, func):
+    def _already_a_wrapper(self, func):
         return getattr(func, self._profiler_wrapped_marker, None) == id(self)
 
-    def _mark_wrapped(self, func):
-        setattr(func, self._profiler_wrapped_marker, id(self))
-        return func
+    def _mark_wrapper(self, wrapper):
+        setattr(wrapper, self._profiler_wrapped_marker, id(self))
+        return wrapper
 
     def run(self, cmd):
         """ Profile a single executable statment in the main namespace.