Create trace of your python program.

pytrace.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2013 by Xu Cheng ([email protected])
#

import sys
if sys.version[0] == '2':
    from io import open

import cgi
import glob
import os
import optparse
import shutil
import trace

if sys.version[0] == '2':
    def print_(*args, **kwargs):
        fp = kwargs.pop("file", sys.stdout)
        if fp is None:
            return

        def write(data):
            if not isinstance(data, basestring):
                data = str(data)
            fp.write(data)
        want_unicode = False
        sep = kwargs.pop("sep", None)
        if sep is not None:
            if isinstance(sep, unicode):
                want_unicode = True
            elif not isinstance(sep, str):
                raise TypeError("sep must be None or a string")
        end = kwargs.pop("end", None)
        if end is not None:
            if isinstance(end, unicode):
                want_unicode = True
            elif not isinstance(end, str):
                raise TypeError("end must be None or a string")
        if kwargs:
            raise TypeError("invalid keyword arguments to print()")
        if not want_unicode:
            for arg in args:
                if isinstance(arg, unicode):
                    want_unicode = True
                    break
        if want_unicode:
            newline = unicode("\n")
            space = unicode(" ")
        else:
            newline = "\n"
            space = " "
        if sep is None:
            sep = space
        if end is None:
            end = newline
        for i, arg in enumerate(args):
            if i:
                write(sep)
            write(arg)
        write(end)
else:
    import builtins
    print_ = getattr(builtins, "print")
    del builtins


def gen_coverage(output_dir, script, argv=None, working_dir=None):
    if not working_dir:
        working_dir = os.path.dirname(os.path.realpath(script))
    script = os.path.realpath(script)
    old_current_dir = os.getcwd()
    old_sys_argv = sys.argv[:]
    old_sys_path = sys.path[:]
    os.chdir(working_dir)
    try:
        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)
        os.makedirs(output_dir)

        print_('>>> Running the script:', script)
        print_('    with working directory:', working_dir)

        with open(script, mode='r', encoding='utf-8') as f:
            if sys.version[0] == '2':
                code = f.read().encode('ascii')
            else:
                code = f.read()

        compiled_code = compile(code, script, 'exec')
        globs = {
            '__file__': script,
            '__name__': '__main__',
            '__package__': None,
            '__cached__': None,
        }
        sys.argv = [script]
        sys.path.insert(0, working_dir)
        if argv:
            sys.argv.extend(argv)
        tracer = trace.Trace(trace=0, count=1,
                             ignoredirs=[sys.prefix, sys.exec_prefix],
                             outfile=os.path.join(output_dir, u'count'))
        try:
            tracer.runctx(compiled_code, globs, globs)
        except:
            pass

        print_('>>> Generate coverage files')
        res = tracer.results()
        res.write_results(show_missing=True, coverdir=output_dir)
    finally:
        os.chdir(old_current_dir)
        sys.argv = old_sys_argv
        sys.path = old_sys_path


def gen_html(input_dir, output_dir):
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir)

    print_('>>> Generate html files')

    # generate css
    with open(os.path.join(output_dir, u'style.css'), mode='w+', encoding='utf-8') as f:
        f.write(u'''
body {
    font-family:"Trebuchet MS",Verdana,"DejaVuSans","VeraSans",Arial,Helvetica,sans-serif;
    font-size: 10pt;
    background-color: white;
    }
.noncovered { background-color:#ffcaca; }
.covered { }
td,th { padding-left:5px;
        padding-right:5px;
        border: 1px solid #ccc;
        font-family:"DejaVu Sans Mono","Bitstream Vera Sans Mono",monospace;
        font-size: 8pt;
   }
th { font-weight:bold; background-color:#eee;}
table { border-collapse: collapse; }
    ''')

    index_html = u''
    # convert each .cover to .html
    for filename in glob.glob(os.path.join(input_dir, '*.cover')):
        print_('  Processing', filename)
        html_table = u'<table><thead><th>Run count</th><th>Line number</th><th>Code</th></thead><tbody>'
        with open(filename, mode='r', encoding='utf-8') as filein:
            linecounter = 0
            noncoveredLineCounter = 0
            for line in filein:
                linecounter += 1
                runcount = u''
                if line[5] == u':':
                    runcount = cgi.escape(line[:5].strip())
                cssClass = u'covered'
                if line.startswith(u'>>>>>>'):
                    noncoveredLineCounter += 1
                    cssClass = u'noncovered'
                    runcount = u'&#x25ba;'
                html_table += u'<tr class="%s"><td align="right">%s</td><td align="right">%d</td><td nowrap>%s</td></tr>\n' % (
                    cssClass, runcount, linecounter, cgi.escape(line[7:].rstrip()).replace(u' ', u'&nbsp;'))
        html_table += u'</tbody></table>'
        if filename.startswith(input_dir):
            source_file_name = filename[len(input_dir):]
        else:
            source_file_name = filename
        source_file_name = source_file_name.lstrip(u'/\\')[:-6] + u'.py'
        coverage_percent = int(
            100 * float(linecounter - noncoveredLineCounter) / float(linecounter))
        html = u'''<html><!-- Generated by pytrace.py - [email protected] --><head><link rel="stylesheet" href="style.css" type="text/css"></head><body>
<b>File:</b> %s<br>
<b>Coverage:</b> %d%% &nbsp; ( <span class="noncovered">&nbsp;&#x25ba;&nbsp;</span> = Code not executed. )<br>
<br>
        ''' % (cgi.escape(source_file_name), coverage_percent) + html_table + u'</body></html>'
        with open(os.path.join(output_dir, source_file_name + u'.html'), mode='w+', encoding='utf-8') as fileout:
            fileout.write(html)
        index_html += u'<tr><td><a href="%s">%s</a></td><td>%d%%</td></tr>\n' % (
            source_file_name + u'.html', cgi.escape(source_file_name), coverage_percent)

    with open(os.path.join(output_dir, u'index.html'), mode='w+', encoding='utf-8') as f:
        f.write(u'''<html><head><link rel="stylesheet" href="style.css" type="text/css"></head>
<body><table><thead><th>File</th><th>Coverage</th></thead><tbody>%s</tbody></table></body></html>''' % index_html)


if __name__ == '__main__':
    optparse.OptionParser.format_description = lambda self, formatter: self.description
    parser = optparse.OptionParser(
        usage="%prog [options] SCRIPT_FILE [args]", version="%prog 1.0",
        description='''Create trace of your python program.
Copyright (c) 2013 by Xu Cheng ([email protected])
''')
    parser.add_option('-o', '--output', action="store", type="string",
                      dest="output_dir", default=u"./output", help="set output directory")
    parser.add_option('-w', '--working', action="store", type="string",
                      dest="working_dir", help="set working directory when run the script")
    options, args = parser.parse_args()

    if len(args) == 0:
        parser.error("The path of your script is required.")

    if not os.path.exists(args[0]):
        parser.error("The path of your script is not found.")

    if len(args) == 1:
        running_args = None
    else:
        running_args = args[1:]

    output_dir = options.output_dir.decode(
        'utf-8') if sys.version[0] == '2' else options.output_dir
    coverage_dir = os.path.join(output_dir, u'coverage')
    html_dir = os.path.join(output_dir, u'html')
    if options.working_dir:
        if sys.version[0] == '2':
            gen_coverage(
                coverage_dir, args[0], running_args, options.working_dir.decode('utf-8'))
        else:
            gen_coverage(
                coverage_dir, args[0], running_args, options.working_dir)
    else:
        gen_coverage(coverage_dir, args[0], running_args)
    gen_html(coverage_dir, html_dir)