prof.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  1. """This plugin will run tests using the hotshot profiler, which is part
  2. of the standard library. To turn it on, use the ``--with-profile`` option
  3. or set the NOSE_WITH_PROFILE environment variable. Profiler output can be
  4. controlled with the ``--profile-sort`` and ``--profile-restrict`` options,
  5. and the profiler output file may be changed with ``--profile-stats-file``.
  6. See the `hotshot documentation`_ in the standard library documentation for
  7. more details on the various output options.
  8. .. _hotshot documentation: http://docs.python.org/library/hotshot.html
  9. """
  10. try:
  11. import hotshot
  12. from hotshot import stats
  13. except ImportError:
  14. hotshot, stats = None, None
  15. import logging
  16. import os
  17. import sys
  18. import tempfile
  19. from nose.plugins.base import Plugin
  20. from nose.util import tolist
  21. log = logging.getLogger('nose.plugins')
  22. class Profile(Plugin):
  23. """
  24. Use this plugin to run tests using the hotshot profiler.
  25. """
  26. pfile = None
  27. clean_stats_file = False
  28. def options(self, parser, env):
  29. """Register commandline options.
  30. """
  31. if not self.available():
  32. return
  33. Plugin.options(self, parser, env)
  34. parser.add_option('--profile-sort', action='store', dest='profile_sort',
  35. default=env.get('NOSE_PROFILE_SORT', 'cumulative'),
  36. metavar="SORT",
  37. help="Set sort order for profiler output")
  38. parser.add_option('--profile-stats-file', action='store',
  39. dest='profile_stats_file',
  40. metavar="FILE",
  41. default=env.get('NOSE_PROFILE_STATS_FILE'),
  42. help='Profiler stats file; default is a new '
  43. 'temp file on each run')
  44. parser.add_option('--profile-restrict', action='append',
  45. dest='profile_restrict',
  46. metavar="RESTRICT",
  47. default=env.get('NOSE_PROFILE_RESTRICT'),
  48. help="Restrict profiler output. See help for "
  49. "pstats.Stats for details")
  50. def available(cls):
  51. return hotshot is not None
  52. available = classmethod(available)
  53. def begin(self):
  54. """Create profile stats file and load profiler.
  55. """
  56. if not self.available():
  57. return
  58. self._create_pfile()
  59. self.prof = hotshot.Profile(self.pfile)
  60. def configure(self, options, conf):
  61. """Configure plugin.
  62. """
  63. if not self.available():
  64. self.enabled = False
  65. return
  66. Plugin.configure(self, options, conf)
  67. self.conf = conf
  68. if options.profile_stats_file:
  69. self.pfile = options.profile_stats_file
  70. self.clean_stats_file = False
  71. else:
  72. self.pfile = None
  73. self.clean_stats_file = True
  74. self.fileno = None
  75. self.sort = options.profile_sort
  76. self.restrict = tolist(options.profile_restrict)
  77. def prepareTest(self, test):
  78. """Wrap entire test run in :func:`prof.runcall`.
  79. """
  80. if not self.available():
  81. return
  82. log.debug('preparing test %s' % test)
  83. def run_and_profile(result, prof=self.prof, test=test):
  84. self._create_pfile()
  85. prof.runcall(test, result)
  86. return run_and_profile
  87. def report(self, stream):
  88. """Output profiler report.
  89. """
  90. log.debug('printing profiler report')
  91. self.prof.close()
  92. prof_stats = stats.load(self.pfile)
  93. prof_stats.sort_stats(self.sort)
  94. # 2.5 has completely different stream handling from 2.4 and earlier.
  95. # Before 2.5, stats objects have no stream attribute; in 2.5 and later
  96. # a reference sys.stdout is stored before we can tweak it.
  97. compat_25 = hasattr(prof_stats, 'stream')
  98. if compat_25:
  99. tmp = prof_stats.stream
  100. prof_stats.stream = stream
  101. else:
  102. tmp = sys.stdout
  103. sys.stdout = stream
  104. try:
  105. if self.restrict:
  106. log.debug('setting profiler restriction to %s', self.restrict)
  107. prof_stats.print_stats(*self.restrict)
  108. else:
  109. prof_stats.print_stats()
  110. finally:
  111. if compat_25:
  112. prof_stats.stream = tmp
  113. else:
  114. sys.stdout = tmp
  115. def finalize(self, result):
  116. """Clean up stats file, if configured to do so.
  117. """
  118. if not self.available():
  119. return
  120. try:
  121. self.prof.close()
  122. except AttributeError:
  123. # TODO: is this trying to catch just the case where not
  124. # hasattr(self.prof, "close")? If so, the function call should be
  125. # moved out of the try: suite.
  126. pass
  127. if self.clean_stats_file:
  128. if self.fileno:
  129. try:
  130. os.close(self.fileno)
  131. except OSError:
  132. pass
  133. try:
  134. os.unlink(self.pfile)
  135. except OSError:
  136. pass
  137. return None
  138. def _create_pfile(self):
  139. if not self.pfile:
  140. self.fileno, self.pfile = tempfile.mkstemp()
  141. self.clean_stats_file = True