Rev 4643: lsprof support. in http://bazaar.launchpad.net/~lifeless/bzr/test-speed

[Robert Collins]

At http://bazaar.launchpad.net/~lifeless/bzr/test-speed

------------------------------------------------------------
revno: 4643
revision-id: robertc at robertcollins.net-20090824210509-pproia2q9evq1nsl
parent: robertc at robertcollins.net-20090824202825-dlt301yil857ksuc
committer: Robert Collins <robertc at robertcollins.net>
branch nick: test-speed
timestamp: Tue 2009-08-25 07:05:09 +1000
message:
  lsprof support.
=== modified file 'NEWS'
--- a/NEWS	2009-08-24 18:28:46 +0000
+++ b/NEWS	2009-08-24 21:05:09 +0000
@@ -67,6 +67,10 @@
 Internals
 *********
 
+* The ``bzrlib.lsprof`` module has a new class ``BzrProfiler`` which makes
+  profiling in some situations like callbacks and generators easier.
+  (Robert Collins)
+
 Testing
 *******
 

=== modified file 'bzrlib/lsprof.py'
--- a/bzrlib/lsprof.py	2009-03-08 06:18:06 +0000
+++ b/bzrlib/lsprof.py	2009-08-24 21:05:09 +0000
@@ -13,45 +13,74 @@
 
 __all__ = ['profile', 'Stats']
 
-_g_threadmap = {}
-
-
-def _thread_profile(f, *args, **kwds):
-    # we lose the first profile point for a new thread in order to trampoline
-    # a new Profile object into place
-    global _g_threadmap
-    thr = thread.get_ident()
-    _g_threadmap[thr] = p = Profiler()
-    # this overrides our sys.setprofile hook:
-    p.enable(subcalls=True, builtins=True)
-
-
 def profile(f, *args, **kwds):
     """Run a function profile.
 
     Exceptions are not caught: If you need stats even when exceptions are to be
-    raised, passing in a closure that will catch the exceptions and transform
-    them appropriately for your driver function.
+    raised, pass in a closure that will catch the exceptions and transform them
+    appropriately for your driver function.
 
     :return: The functions return value and a stats object.
     """
-    global _g_threadmap
-    p = Profiler()
-    p.enable(subcalls=True)
-    threading.setprofile(_thread_profile)
+    profiler = BzrProfiler()
+    profiler.start()
     try:
         ret = f(*args, **kwds)
     finally:
-        p.disable()
-        for pp in _g_threadmap.values():
+        stats = profiler.stop()
+    return ret, stats
+
+
+class BzrProfiler(object):
+    """Bzr utility wrapper around Profiler.
+    
+    For most uses the module level 'profile()' function will be suitable.
+    However profiling when a simple wrapped function isn't available may
+    be easier to accomplish using this class.
+
+    To use it, create a BzrProfiler and call start() on it. Some arbitrary
+    time later call stop() to stop profiling and retrieve the statistics
+    from the code executed in the interim.
+    """
+
+    def start(self):
+        """Start profiling.
+        
+        This hooks into threading and will record all calls made until
+        stop() is called.
+        """
+        self._g_threadmap = {}
+        self.p = Profiler()
+        self.p.enable(subcalls=True)
+        threading.setprofile(self._thread_profile)
+
+    def stop(self):
+        """Stop profiling.
+
+        This unhooks from threading and cleans up the profiler, returning
+        the gathered Stats object.
+
+        :return: A bzrlib.lsprof.Stats object.
+        """
+        self.p.disable()
+        for pp in self._g_threadmap.values():
             pp.disable()
         threading.setprofile(None)
+        p = self.p
+        self.p = None
+        threads = {}
+        for tid, pp in self._g_threadmap.items():
+            threads[tid] = Stats(pp.getstats(), {})
+        self._g_threadmap = None
+        return Stats(p.getstats(), threads)
 
-    threads = {}
-    for tid, pp in _g_threadmap.items():
-        threads[tid] = Stats(pp.getstats(), {})
-    _g_threadmap = {}
-    return ret, Stats(p.getstats(), threads)
+    def _thread_profile(self, f, *args, **kwds):
+        # we lose the first profile point for a new thread in order to
+        # trampoline a new Profile object into place
+        thr = thread.get_ident()
+        self._g_threadmap[thr] = p = Profiler()
+        # this overrides our sys.setprofile hook:
+        p.enable(subcalls=True, builtins=True)
 
 
 class Stats(object):

=== modified file 'bzrlib/tests/test_lsprof.py'
--- a/bzrlib/tests/test_lsprof.py	2009-03-23 14:59:43 +0000
+++ b/bzrlib/tests/test_lsprof.py	2009-08-24 21:05:09 +0000
@@ -92,3 +92,22 @@
         self.stats.save(f)
         data1 = cPickle.load(open(f))
         self.assertEqual(type(data1), bzrlib.lsprof.Stats)
+
+
+class TestBzrProfiler(tests.TestCase):
+
+    _test_needs_features = [LSProfFeature]
+
+    def test_start_call_stuff_stop(self):
+        profiler = bzrlib.lsprof.BzrProfiler()
+        profiler.start()
+        try:
+            def a_function():
+                pass
+            a_function()
+        finally:
+            stats = profiler.stop()
+        stats.freeze()
+        lines = [str(data) for data in stats.data]
+        lines = [line for line in lines if 'a_function' in line]
+        self.assertLength(1, lines)