profilehooks.py

Go to the documentation of this file.

##
# 
# Profiling hooks
# 
# This module contains a couple of decorators (`profile` and `coverage`) that
# can be used to wrap functions and/or methods to produce profiles and line
# coverage reports.
# 
# Usage example (Python 2.4 or newer)::
# 
#     from profilehooks import profile, coverage
# 
#     @profile    # or @coverage
#     def fn(n):
#         if n < 2: return 1
#         else: return n * fn(n-1)
# 
#     print fn(42)
# 
# Usage example (Python 2.3 or older)::
# 
#     from profilehooks import profile, coverage
# 
#     def fn(n):
#         if n < 2: return 1
#         else: return n * fn(n-1)
# 
#     # Now wrap that function in a decorator
#     fn = profile(fn) # or coverage(fn)
# 
#     print fn(42)
# 
# Reports for all thusly decorated functions will be printed to sys.stdout
# on program termination.  You can alternatively request for immediate
# reports for each call by passing immediate=True to the profile decorator.
# 
# There's also a @timecall decorator for printing the time to sys.stderr
# every time a function is called, when you just want to get a rough measure
# instead of a detailed (but costly) profile.
# 
# Caveats
# 
#   A thread on python-dev convinced me that hotshot produces bogus numbers.
#   See http://mail.python.org/pipermail/python-dev/2005-November/058264.html
# 
#   I don't know what will happen if a decorated function will try to call
#   another decorated function.  All decorators probably need to explicitly
#   support nested profiling (currently TraceFuncCoverage is the only one
#   that supports this, while HotShotFuncProfile has support for recursive
#   functions.)
# 
#   Profiling with hotshot creates temporary files (*.prof for profiling,
#   *.cprof for coverage) in the current directory.  These files are not
#   cleaned up.  Exception: when you specify a filename to the profile
#   decorator (to store the pstats.Stats object for later inspection),
#   the temporary file will be the filename you specified with '.raw'
#   appended at the end.
# 
#   Coverage analysis with hotshot seems to miss some executions resulting
#   in lower line counts and some lines errorneously marked as never
#   executed.  For this reason coverage analysis now uses trace.py which is
#   slower, but more accurate.
# 
# Copyright (c) 2004--2008 Marius Gedminas <marius@pov.lt>
# Copyright (c) 2007 Hanno Schlichting
# Copyright (c) 2008 Florian Schulze
# 
# Released under the MIT licence since December 2006:
# 
#     Permission is hereby granted, free of charge, to any person obtaining a
#     copy of this software and associated documentation files (the "Software"),
#     to deal in the Software without restriction, including without limitation
#     the rights to use, copy, modify, merge, publish, distribute, sublicense,
#     and/or sell copies of the Software, and to permit persons to whom the
#     Software is furnished to do so, subject to the following conditions:
# 
#     The above copyright notice and this permission notice shall be included in
#     all copies or substantial portions of the Software.
# 
#     THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#     IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#     FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#     AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#     LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#     FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#     DEALINGS IN THE SOFTWARE.
# 
# (Previously it was distributed under the GNU General Public Licence.)
# 
# $Id$

__author__ = "Marius Gedminas (marius@gedmin.as)"
__copyright__ = "Copyright 2004-2008, Marius Gedminas"
__license__ = "MIT"
__version__ = "1.3"
__date__ = "2008-06-10"


import atexit
import inspect
import sys
import re

# For profiling
from profile import Profile
import pstats

# For hotshot profiling (inaccurate!)
import hotshot
import hotshot.stats

# For trace.py coverage
import trace

# For hotshot coverage (inaccurate!; uses undocumented APIs; might break)
import _hotshot
import hotshot.log

# For timecall
import time


##
# Mark `fn` for profiling.
# 
#     If `skip` is > 0, first `skip` calls to `fn` will not be profiled.
# 
#     If `immediate` is False, profiling results will be printed to
#     sys.stdout on program termination.  Otherwise results will be printed
#     after each call.
# 
#     If `dirs` is False only the name of the file will be printed.
#     Otherwise the full path is used.
# 
#     `sort` can be a list of sort keys (defaulting to ['cumulative',
#     'time', 'calls']).  The following ones are recognized::
# 
#         'calls'      -- call count
#         'cumulative' -- cumulative time
#         'file'       -- file name
#         'line'       -- line number
#         'module'     -- file name
#         'name'       -- function name
#         'nfl'        -- name/file/line
#         'pcalls'     -- call count
#         'stdname'    -- standard name
#         'time'       -- internal time
# 
#     `entries` limits the output to the first N entries.
# 
#     If `filename` is specified, the profile stats will be stored in the
#     named file.  You can load them pstats.Stats(filename).
# 
#     Usage::
# 
#         def fn(...):
#             ...
#         fn = profile(fn, skip=1)
# 
#     If you are using Python 2.4, you should be able to use the decorator
#     syntax::
# 
#         @profile(skip=3)
#         def fn(...):
#             ...
# 
#     or just ::
# 
#         @profile
#         def fn(...):
#             ...
# 
#     
def profile(fn=None, skip=0, filename=None, immediate=False, dirs=False,
            sort=None, entries=40):
    if fn is None: # @profile() syntax -- we are a decorator maker
        def decorator(fn):
            return profile(fn, skip=skip, filename=filename,
                           immediate=immediate, dirs=dirs,
                           sort=sort, entries=entries)
        return decorator
    # @profile syntax -- we are a decorator.
    fp = FuncProfile(fn, skip=skip, filename=filename,
                     immediate=immediate, dirs=dirs,
                     sort=sort, entries=entries)
    # fp = HotShotFuncProfile(fn, skip=skip, filename=filename, ...)
         # or HotShotFuncProfile
    # We cannot return fp or fp.__call__ directly as that would break method
    # definitions, instead we need to return a plain function.
    def new_fn(*args, **kw):
        return fp(*args, **kw)
    new_fn.__doc__ = fn.__doc__
    new_fn.__name__ = fn.__name__
    new_fn.__dict__ = fn.__dict__
    new_fn.__module__ = fn.__module__
    return new_fn


##
# Mark `fn` for line coverage analysis.
# 
#     Results will be printed to sys.stdout on program termination.
# 
#     Usage::
# 
#         def fn(...):
#             ...
#         fn = coverage(fn)
# 
#     If you are using Python 2.4, you should be able to use the decorator
#     syntax::
# 
#         @coverage
#         def fn(...):
#             ...
# 
#     
def coverage(fn):
    fp = TraceFuncCoverage(fn) # or HotShotFuncCoverage
    # We cannot return fp or fp.__call__ directly as that would break method
    # definitions, instead we need to return a plain function.
    def new_fn(*args, **kw):
        return fp(*args, **kw)
    new_fn.__doc__ = fn.__doc__
    new_fn.__name__ = fn.__name__
    new_fn.__dict__ = fn.__dict__
    new_fn.__module__ = fn.__module__
    return new_fn


##
# Mark `fn` for line coverage analysis.
# 
#     Uses the 'hotshot' module for fast coverage analysis.
# 
#     BUG: Produces inaccurate results.
# 
#     See the docstring of `coverage` for usage examples.
#     
def coverage_with_hotshot(fn):
    fp = HotShotFuncCoverage(fn)
    # We cannot return fp or fp.__call__ directly as that would break method
    # definitions, instead we need to return a plain function.
    def new_fn(*args, **kw):
        return fp(*args, **kw)
    new_fn.__doc__ = fn.__doc__
    new_fn.__name__ = fn.__name__
    new_fn.__dict__ = fn.__dict__
    new_fn.__module__ = fn.__module__
    return new_fn


##
# Profiler for a function (uses profile).
class FuncProfile:

    # This flag is shared between all instances
    in_profiler = False

    ##
    # Creates a profiler for a function.
    # 
    #         Every profiler has its own log file (the name of which is derived
    #         from the function name).
    # 
    #         FuncProfile registers an atexit handler that prints profiling
    #         information to sys.stderr when the program terminates.
    #         
    def __init__(self, fn, skip=0, filename=None, immediate=False, dirs=False,
                 sort=None, entries=40):
        self.fn = fn
        self.filename = filename
        self.immediate = immediate
        self.dirs = dirs
        self.sort = sort or ('cumulative', 'time', 'calls')
        if isinstance(self.sort, str):
            self.sort = (self.sort, )
        self.entries = entries
        self.stats = pstats.Stats(Profile())
        self.ncalls = 0
        self.skip = skip
        self.skipped = 0
        atexit.register(self.atexit)

    ##
    # Profile a singe call to the function.
    def __call__(self, *args, **kw):
        self.ncalls += 1
        if self.skip > 0:
            self.skip -= 1
            self.skipped += 1
            return self.fn(*args, **kw)
        if FuncProfile.in_profiler:
            # handle recursive calls
            return self.fn(*args, **kw)
        # You cannot reuse the same profiler for many calls and accumulate
        # stats that way.  :-/
        profiler = Profile()
        try:
            FuncProfile.in_profiler = True
            return profiler.runcall(self.fn, *args, **kw)
        finally:
            FuncProfile.in_profiler = False
            self.stats.add(profiler)
            if self.immediate:
                self.print_stats()
                self.reset_stats()

    ##
    # Print profile information to sys.stdout.
    def print_stats(self):
        funcname = self.fn.__name__
        filename = self.fn.func_code.co_filename
        lineno = self.fn.func_code.co_firstlineno
        print
        print "*** PROFILER RESULTS ***"
        print "%s (%s:%s)" % (funcname, filename, lineno)
        print "function called %d times" % self.ncalls,
        if self.skipped:
            print "(%d calls not profiled)" % self.skipped
        else:
            print
        print
        stats = self.stats
        if self.filename:
            stats.dump_stats(self.filename)
        if not self.dirs:
            stats.strip_dirs()
        stats.sort_stats(*self.sort)
        stats.print_stats(self.entries)

    ##
    # Reset accumulated profiler statistics.
    def reset_stats(self):
        self.stats = pstats.Stats(Profile())
        self.ncalls = 0
        self.skipped = 0

    ##
    # Stop profiling and print profile information to sys.stdout.
    # 
    #         This function is registered as an atexit hook.
    #         
    def atexit(self):
        if not self.immediate:
            self.print_stats()


##
# Profiler for a function (uses hotshot).
class HotShotFuncProfile:

    # This flag is shared between all instances
    in_profiler = False

    ##
    # Creates a profiler for a function.
    # 
    #         Every profiler has its own log file (the name of which is derived
    #         from the function name).
    # 
    #         HotShotFuncProfile registers an atexit handler that prints
    #         profiling information to sys.stderr when the program terminates.
    # 
    #         The log file is not removed and remains there to clutter the
    #         current working directory.
    #         
    def __init__(self, fn, skip=0, filename=None):
        self.fn = fn
        self.filename = filename
        if self.filename:
            self.logfilename = filename + ".raw"
        else:
            self.logfilename = fn.__name__ + ".prof"
        self.profiler = hotshot.Profile(self.logfilename)
        self.ncalls = 0
        self.skip = skip
        self.skipped = 0
        atexit.register(self.atexit)

    ##
    # Profile a singe call to the function.
    def __call__(self, *args, **kw):
        self.ncalls += 1
        if self.skip > 0:
            self.skip -= 1
            self.skipped += 1
            return self.fn(*args, **kw)
        if HotShotFuncProfile.in_profiler:
            # handle recursive calls
            return self.fn(*args, **kw)
        try:
            HotShotFuncProfile.in_profiler = True
            return self.profiler.runcall(self.fn, *args, **kw)
        finally:
            HotShotFuncProfile.in_profiler = False

    ##
    # Stop profiling and print profile information to sys.stderr.
    # 
    #         This function is registered as an atexit hook.
    #         
    def atexit(self):
        self.profiler.close()
        funcname = self.fn.__name__
        filename = self.fn.func_code.co_filename
        lineno = self.fn.func_code.co_firstlineno
        print
        print "*** PROFILER RESULTS ***"
        print "%s (%s:%s)" % (funcname, filename, lineno)
        print "function called %d times" % self.ncalls,
        if self.skipped:
            print "(%d calls not profiled)" % self.skipped
        else:
            print
        print
        stats = hotshot.stats.load(self.logfilename)
        # hotshot.stats.load takes ages, and the .prof file eats megabytes, but
        # a saved stats object is small and fast
        if self.filename:
            stats.dump_stats(self.filename)
        # it is best to save before strip_dirs
        stats.strip_dirs()
        stats.sort_stats('cumulative', 'time', 'calls')
        stats.print_stats(40)


##
# Coverage analysis for a function (uses _hotshot).
# 
#     HotShot coverage is reportedly faster than trace.py, but it appears to
#     have problems with exceptions; also line counts in coverage reports
#     are generally lower from line counts produced by TraceFuncCoverage.
#     Is this my bug, or is it a problem with _hotshot?
#     
class HotShotFuncCoverage:

    ##
    # Creates a profiler for a function.
    # 
    #         Every profiler has its own log file (the name of which is derived
    #         from the function name).
    # 
    #         HotShotFuncCoverage registers an atexit handler that prints
    #         profiling information to sys.stderr when the program terminates.
    # 
    #         The log file is not removed and remains there to clutter the
    #         current working directory.
    #         
    def __init__(self, fn):
        self.fn = fn
        self.logfilename = fn.__name__ + ".cprof"
        self.profiler = _hotshot.coverage(self.logfilename)
        self.ncalls = 0
        atexit.register(self.atexit)

    ##
    # Profile a singe call to the function.
    def __call__(self, *args, **kw):
        self.ncalls += 1
        return self.profiler.runcall(self.fn, args, kw)

    ##
    # Stop profiling and print profile information to sys.stderr.
    # 
    #         This function is registered as an atexit hook.
    #         
    def atexit(self):
        self.profiler.close()
        funcname = self.fn.__name__
        filename = self.fn.func_code.co_filename
        lineno = self.fn.func_code.co_firstlineno
        print
        print "*** COVERAGE RESULTS ***"
        print "%s (%s:%s)" % (funcname, filename, lineno)
        print "function called %d times" % self.ncalls
        print
        fs = FuncSource(self.fn)
        reader = hotshot.log.LogReader(self.logfilename)
        for what, (filename, lineno, funcname), tdelta in reader:
            if filename != fs.filename:
                continue
            if what == hotshot.log.LINE:
                fs.mark(lineno)
            if what == hotshot.log.ENTER:
                # hotshot gives us the line number of the function definition
                # and never gives us a LINE event for the first statement in
                # a function, so if we didn't perform this mapping, the first
                # statement would be marked as never executed
                if lineno == fs.firstlineno:
                    lineno = fs.firstcodelineno
                fs.mark(lineno)
        reader.close()
        print fs


##
# Coverage analysis for a function (uses trace module).
# 
#     HotShot coverage analysis is reportedly faster, but it appears to have
#     problems with exceptions.
#     
class TraceFuncCoverage:

    # Shared between all instances so that nested calls work
    tracer = trace.Trace(count=True, trace=False,
                         ignoredirs=[sys.prefix, sys.exec_prefix])

    # This flag is also shared between all instances
    tracing = False

    ##
    # Creates a profiler for a function.
    # 
    #         Every profiler has its own log file (the name of which is derived
    #         from the function name).
    # 
    #         TraceFuncCoverage registers an atexit handler that prints
    #         profiling information to sys.stderr when the program terminates.
    # 
    #         The log file is not removed and remains there to clutter the
    #         current working directory.
    #         
    def __init__(self, fn):
        self.fn = fn
        self.logfilename = fn.__name__ + ".cprof"
        self.ncalls = 0
        atexit.register(self.atexit)

    ##
    # Profile a singe call to the function.
    def __call__(self, *args, **kw):
        self.ncalls += 1
        if TraceFuncCoverage.tracing:
            return self.fn(*args, **kw)
        try:
            TraceFuncCoverage.tracing = True
            return self.tracer.runfunc(self.fn, *args, **kw)
        finally:
            TraceFuncCoverage.tracing = False

    ##
    # Stop profiling and print profile information to sys.stderr.
    # 
    #         This function is registered as an atexit hook.
    #         
    def atexit(self):
        funcname = self.fn.__name__
        filename = self.fn.func_code.co_filename
        lineno = self.fn.func_code.co_firstlineno
        print
        print "*** COVERAGE RESULTS ***"
        print "%s (%s:%s)" % (funcname, filename, lineno)
        print "function called %d times" % self.ncalls
        print
        fs = FuncSource(self.fn)
        for (filename, lineno), count in self.tracer.counts.items():
            if filename != fs.filename:
                continue
            fs.mark(lineno, count)
        print fs
        never_executed = fs.count_never_executed()
        if never_executed:
            print "%d lines were not executed." % never_executed


##
# Source code annotator for a function.
class FuncSource:

    blank_rx = re.compile(r"^\s*finally:\s*(#.*)?$")

    def __init__(self, fn):
        self.fn = fn
        self.filename = inspect.getsourcefile(fn)
        self.source, self.firstlineno = inspect.getsourcelines(fn)
        self.sourcelines = {}
        self.firstcodelineno = self.firstlineno
        self.find_source_lines()

    ##
    # Mark all executable source lines in fn as executed 0 times.
    def find_source_lines(self):
        strs = trace.find_strings(self.filename)
        lines = trace.find_lines_from_code(self.fn.func_code, strs)
        self.firstcodelineno = sys.maxint
        for lineno in lines:
            self.firstcodelineno = min(self.firstcodelineno, lineno)
            self.sourcelines.setdefault(lineno, 0)
        if self.firstcodelineno == sys.maxint:
            self.firstcodelineno = self.firstlineno

    ##
    # Mark a given source line as executed count times.
    # 
    #         Multiple calls to mark for the same lineno add up.
    #         
    def mark(self, lineno, count=1):
        self.sourcelines[lineno] = self.sourcelines.get(lineno, 0) + count

    ##
    # Count statements that were never executed.
    def count_never_executed(self):
        lineno = self.firstlineno
        counter = 0
        for line in self.source:
            if self.sourcelines.get(lineno) == 0:
                if not self.blank_rx.match(line):
                    counter += 1
            lineno += 1
        return counter

    ##
    # Return annotated source code for the function.
    def __str__(self):
        lines = []
        lineno = self.firstlineno
        for line in self.source:
            counter = self.sourcelines.get(lineno)
            if counter is None:
                prefix = ' ' * 7
            elif counter == 0:
                if self.blank_rx.match(line):
                    prefix = ' ' * 7
                else:
                    prefix = '>' * 6 + ' '
            else:
                prefix = '%5d: ' % counter
            lines.append(prefix + line)
            lineno += 1
        return ''.join(lines)


##
# Wrap `fn` and print its execution time.
# 
#     Example::
# 
#         @timecall
#         def somefunc(x, y):
#             time.sleep(x * y)
# 
#         somefunc(2, 3)
# 
#     will print the time taken by somefunc on every call.  If you want just
#     a summary at program termination, use
# 
#         @timecall(immediate=False)
# 
#     
def timecall(fn=None, immediate=True):
    if fn is None: # @timecall() syntax -- we are a decorator maker
        def decorator(fn):
            return timecall(fn, immediate=immediate)
        return decorator
    # @timecall syntax -- we are a decorator.
    fp = FuncTimer(fn, immediate=immediate)
    # We cannot return fp or fp.__call__ directly as that would break method
    # definitions, instead we need to return a plain function.
    def new_fn(*args, **kw):
        return fp(*args, **kw)
    new_fn.__doc__ = fn.__doc__
    new_fn.__name__ = fn.__name__
    new_fn.__dict__ = fn.__dict__
    new_fn.__module__ = fn.__module__
    return new_fn


class FuncTimer(object):

    def __init__(self, fn, immediate):
        self.fn = fn
        self.ncalls = 0
        self.totaltime = 0
        self.immediate = immediate
        if not immediate:
            atexit.register(self.atexit)

    ##
    # Profile a singe call to the function.
    def __call__(self, *args, **kw):
        fn = self.fn
        self.ncalls += 1
        try:
            start = time.time()
            return fn(*args, **kw)
        finally:
            duration = time.time() - start
            self.totaltime += duration
            if self.immediate:
                funcname = fn.__name__
                filename = fn.func_code.co_filename
                lineno = fn.func_code.co_firstlineno
                print >> sys.stderr, "\n  %s (%s:%s):\n    %.3f seconds\n" % (
                                        funcname, filename, lineno, duration)
    def atexit(self):
        if not self.ncalls:
            return
        funcname = self.fn.__name__
        filename = self.fn.func_code.co_filename
        lineno = self.fn.func_code.co_firstlineno
        print ("\n  %s (%s:%s):\n"
               "    %d calls, %.3f seconds (%.3f seconds per call)\n" % (
                                funcname, filename, lineno, self.ncalls,
                                self.totaltime, self.totaltime / self.ncalls))