sqlalchemy/test/lib/profiling.py

"""Profiling support for unit and performance tests.

These are special purpose profiling methods which operate
in a more fine-grained way than nose's profiling plugin.

"""

import os, sys
from test.lib.util import gc_collect, decorator
from nose import SkipTest

__all__ = 'profiled', 'function_call_count', 'conditional_call_count'

all_targets = set()
profile_config = { 'targets': set(),
                   'report': True,
                   'print_callers':False,
                   'print_callees':False,
                   'graphic':False,
                   'sort': ('time', 'calls'),
                   'limit': None }
profiler = None

def profiled(target=None, **target_opts):
    """Optional function profiling.

    @profiled('label')
    or
    @profiled('label', report=True, sort=('calls',), limit=20)

    Enables profiling for a function when 'label' is targetted for
    profiling.  Report options can be supplied, and override the global
    configuration and command-line options.
    """

    # manual or automatic namespacing by module would remove conflict issues
    if target is None:
        target = 'anonymous_target'
    elif target in all_targets:
        print "Warning: redefining profile target '%s'" % target
    all_targets.add(target)

    filename = "%s.prof" % target

    @decorator
    def decorate(fn, *args, **kw):
        if (target not in profile_config['targets'] and
            not target_opts.get('always', None)):
            return fn(*args, **kw)

        elapsed, load_stats, result = _profile(
            filename, fn, *args, **kw)

        graphic = target_opts.get('graphic', profile_config['graphic'])
        if graphic:
            os.system("runsnake %s" % filename)
        else:
            report = target_opts.get('report', profile_config['report'])
            if report:
                sort_ = target_opts.get('sort', profile_config['sort'])
                limit = target_opts.get('limit', profile_config['limit'])
                print "Profile report for target '%s' (%s)" % (
                    target, filename)

                stats = load_stats()
                stats.sort_stats(*sort_)
                if limit:
                    stats.print_stats(limit)
                else:
                    stats.print_stats()

                print_callers = target_opts.get('print_callers', 
                                                profile_config['print_callers'])
                if print_callers:
                    stats.print_callers()

                print_callees = target_opts.get('print_callees', 
                                                profile_config['print_callees'])
                if print_callees:
                    stats.print_callees()

        os.unlink(filename)
        return result
    return decorate

def function_call_count(count=None, versions={}, variance=0.05):
    """Assert a target for a test case's function call count.

    count
      Optional, general target function call count.

    versions
      Optional, a dictionary of Python version strings to counts,
      for example::

        { '2.5.1': 110,
          '2.5': 100,
          '2.4': 150 }

      The best match for the current running python will be used.
      If none match, 'count' will be used as the fallback.

    variance
      An +/- deviation percentage, defaults to 5%.
    """

    # this could easily dump the profile report if --verbose is in effect

    version_info = list(sys.version_info)
    py_version = '.'.join([str(v) for v in sys.version_info])
    try:
        from sqlalchemy.cprocessors import to_float
        cextension = True
    except ImportError:
        cextension = False

    while version_info:
        version = '.'.join([str(v) for v in version_info])
        if cextension and (version + "+cextension") in versions:
            version += "+cextension"
            count = versions[version]
            break
        elif version in versions:
            count = versions[version]
            break

        version_info.pop()

    if count is None:
        print "Warning: no function call count specified for version: '%s'" % py_version
        return lambda fn: fn

    @decorator
    def decorate(fn, *args, **kw):
        try:
            filename = "%s.prof" % fn.__name__

            elapsed, stat_loader, result = _profile(
                filename, fn, *args, **kw)

            stats = stat_loader()
            calls = stats.total_calls

            stats.sort_stats('calls', 'cumulative')
            stats.print_stats()
            #stats.print_callers()
            deviance = int(count * variance)
            if (calls < (count - deviance) or
                calls > (count + deviance)):
                raise AssertionError(
                    "Function call count %s not within %s%% "
                    "of expected %s. (Python version %s)" % (
                    calls, (variance * 100), count, py_version))

            return result
        finally:
            if os.path.exists(filename):
                os.unlink(filename)
    return decorate

def conditional_call_count(discriminator, categories):
    """Apply a function call count conditionally at runtime.

    Takes two arguments, a callable that returns a key value, and a dict
    mapping key values to a tuple of arguments to function_call_count.

    The callable is not evaluated until the decorated function is actually
    invoked.  If the `discriminator` returns a key not present in the
    `categories` dictionary, no call count assertion is applied.

    Useful for integration tests, where running a named test in isolation may
    have a function count penalty not seen in the full suite, due to lazy
    initialization in the DB-API, SA, etc.
    """
    @decorator
    def decorate(fn, *args, **kw):
        criteria = categories.get(discriminator(), None)
        if criteria is None:
            return fn(*args, **kw)

        rewrapped = function_call_count(*criteria)(fn)
        return rewrapped(*args, **kw)
    return decorate


def _profile(filename, fn, *args, **kw):
    global profiler
    if not profiler:
        if sys.version_info > (2, 5):
            try:
                import cProfile
                profiler = 'cProfile'
            except ImportError:
                pass
        if not profiler:
            try:
                import hotshot
                profiler = 'hotshot'
            except ImportError:
                profiler = 'skip'

    if profiler == 'skip':
        raise SkipTest('Profiling not supported on this platform')
    elif profiler == 'cProfile':
        return _profile_cProfile(filename, fn, *args, **kw)
    else:
        return _profile_hotshot(filename, fn, *args, **kw)

def _profile_cProfile(filename, fn, *args, **kw):
    import cProfile, pstats, time

    load_stats = lambda: pstats.Stats(filename)
    gc_collect()

    began = time.time()
    cProfile.runctx('result = fn(*args, **kw)', globals(), locals(),
                    filename=filename)
    ended = time.time()

    return ended - began, load_stats, locals()['result']

def _profile_hotshot(filename, fn, *args, **kw):
    import hotshot, hotshot.stats, time
    load_stats = lambda: hotshot.stats.load(filename)

    gc_collect()
    prof = hotshot.Profile(filename)
    began = time.time()
    prof.start()
    try:
        result = fn(*args, **kw)
    finally:
        prof.stop()
        ended = time.time()
        prof.close()

    return ended - began, load_stats, result