Replace hotshot with cProfile

nose2/plugins/prof.py

@@ -1,8 +1,8 @@
 """
-Profile test execution using hotshot.
+Profile test execution using cProfile.
 
 This plugin implements :func:`startTestRun` and replaces
-``event.executeTests`` with :meth:`hotshot.Profile.runcall`. It
+``event.executeTests`` with :meth:`cProfile.Profile.runcall`. It
 implements :func:`beforeSummaryReport` to output profiling information
 before the final test summary time. Config file options ``filename``,
 ``sort`` and ``restrict`` can be used to change where profiling
@@ -15,14 +15,9 @@
 sections of the configuration file.
 
 """
-try:
-    import hotshot
-    from hotshot import stats
-except ImportError:
-    hotshot, stats = None, None
+import cProfile
+import pstats
 import logging
-import os
-import tempfile
 
 from nose2 import events, util
 
@@ -41,20 +36,11 @@ def __init__(self):
         self.pfile = self.config.as_str('filename', '')
         self.sort = self.config.as_str('sort', 'cumulative')
         self.restrict = self.config.as_list('restrict', [])
-        self.clean = False
         self.fileno = None
 
-    def register(self):
-        """Don't register if hotshot is not found"""
-        if hotshot is None:
-            log.error("Unable to profile: hotshot module not available")
-            return
-        super(Profiler, self).register()
-
     def startTestRun(self, event):
         """Set up the profiler"""
-        self.createPfile()
-        self.prof = hotshot.Profile(self.pfile)
+        self.prof = cProfile.Profile()
         event.executeTests = self.prof.runcall
 
     def beforeSummaryReport(self, event):
@@ -68,34 +54,16 @@ def write(self, *msg):
                     event.stream.write(' ')
                 event.stream.flush()
         stream = Stream()
-        self.prof.close()
-        prof_stats = stats.load(self.pfile)
+        prof_stats = pstats.Stats(self.prof, stream=stream)
         prof_stats.sort_stats(self.sort)
         event.stream.writeln(util.ln("Profiling results"))
-        tmp = prof_stats.stream
-        prof_stats.stream = stream
-        try:
-            if self.restrict:
-                prof_stats.print_stats(*self.restrict)
-            else:
-                prof_stats.print_stats()
-        finally:
-            prof_stats.stream = tmp
-        self.prof.close()
-        event.stream.writeln('')
+        if self.restrict:
+            prof_stats.print_stats(*self.restrict)
+        else:
+            prof_stats.print_stats()
 
-        if self.clean:
-            if self.fileno:
-                try:
-                    os.close(self.fileno)
-                except OSError:
-                    pass
-            try:
-                os.unlink(self.pfile)
-            except OSError:
-                pass
+        if self.pfile:
+            prof_stats.dump_stats(self.pfile)
 
-    def createPfile(self):
-        if not self.pfile:
-            self.fileno, self.pfile = tempfile.mkstemp()
-            self.clean = True
+        self.prof.disable()
+        event.stream.writeln('')

nose2/tests/unit/test_prof_plugin.py

@@ -9,19 +9,20 @@ class TestProfPlugin(TestCase):
 
     def setUp(self):
         self.plugin = prof.Profiler(session=session.Session())
-        self.hotshot = prof.hotshot
-        self.stats = prof.stats
-        prof.hotshot = Stub()
-        prof.stats = Stub()
+        # stub out and save the cProfile and pstats modules
+        self.cProfile = prof.cProfile
+        self.pstats = prof.pstats
+        prof.cProfile = Stub()
+        prof.pstats = Stub()
 
     def tearDown(self):
-        prof.hotshot = self.hotshot
-        prof.stats = self.stats
+        prof.cProfile = self.cProfile
+        prof.pstats = self.pstats
 
     def test_startTestRun_sets_executeTests(self):
         _prof = Stub()
         _prof.runcall = object()
-        prof.hotshot.Profile = lambda filename: _prof
+        prof.cProfile.Profile = lambda: _prof
         event = StartTestRunEvent(runner=None, suite=None, result=None,
                                   startTime=None, executeTests=None)
         self.plugin.startTestRun(event)