Merge rkern#76 from iddl/line_profiler

Feature: Add 'profile all' option

Author: ahippo

_line_profiler.pyx

@@ -109,13 +109,15 @@ cdef class LineProfiler:
     cdef public dict last_time
     cdef public double timer_unit
     cdef public long enable_count
+    cdef public bint profile_all
 
     def __init__(self, *functions):
         self.functions = []
         self.code_map = {}
         self.last_time = {}
         self.timer_unit = hpTimerUnit()
         self.enable_count = 0
+        self.profile_all = False
         for func in functions:
             self.add_function(func)
 
@@ -132,6 +134,16 @@ cdef class LineProfiler:
             self.code_map[code] = {}
             self.functions.append(func)
 
+    def enable_profile_all(self):
+        """ Record line profiling information for all executed code.
+        """
+        self.profile_all = True
+
+    def disable_profile_all(self):
+        """ Disable recording line profiling information for all executed code.
+        """
+        self.profile_all = False
+
     def enable_by_count(self):
         """ Enable the profiler if it hasn't been enabled before.
         """
@@ -198,6 +210,12 @@ cdef int python_trace_callback(object self_, PyFrameObject *py_frame, int what,
     self = <LineProfiler>self_
     last_time = self.last_time
 
+    if self.profile_all and what == PyTrace_CALL and <object>py_frame.f_code not in self.code_map:
+        # enable recording profiling information for this function if
+        # profile_all is enabled
+        code = <object>py_frame.f_code
+        self.code_map[code] = {}
+
     if what == PyTrace_LINE or what == PyTrace_RETURN:
         code = <object>py_frame.f_code
         if code in self.code_map:

tests/test_line_profiler.py

@@ -60,6 +60,18 @@ def test_enable_disable(self):
         self.assertEqual(lp.enable_count, 0)
         self.assertEqual(lp.last_time, {})
 
+    def test_enable_profile_all(self):
+        lp = LineProfiler()
+
+        lp.enable_profile_all()
+        lp.enable()
+        value = f(10)
+        lp.disable()
+
+        self.assertEqual(value, f(10))
+        self.assertEqual(len(lp.code_map.keys()), 1)
+        self.assertEqual(len(lp.code_map[f.__code__]), 2)
+
     def test_function_decorator(self):
         profile = LineProfiler()
         f_wrapped = profile(f)