IPython cell magic for the line_profiler

line_profiler_cell_magic.py

from __future__ import absolute_import, division, print_function

import inspect
import linecache
import os

from IPython.core import magic
from IPython.core import page
from line_profiler import LineProfiler
from six.moves import StringIO


@magic.magics_class
class CPUProfiler(magic.Magics):
    '''IPython cell magic interface to line_profiler.

    Examples
    --------
    >>> %%lprun -f func1 -f func2
    >>> <statements>
    '''

    @magic.cell_magic
    def lprun(self, line='', cell=None):
        functions = []
        words = [word.strip() for word in line.split()]
        for i, word in enumerate(words):
            if word == '-f':
                functions.append(eval(words[i+1], self.shell.user_ns))

        local_ns = {}
        encapsulated_cell = 'def __main__():\n'
        for line in cell.split('\n'):
            encapsulated_cell += '    {}\n'.format(line)
        exec(encapsulated_cell, self.shell.user_ns, local_ns)
        self.shell.user_ns.update(local_ns)
        main = self.shell.user_ns.pop('__main__')

        profiler = LineProfiler()
        for function in functions:
            profiler.add_function(function)
        profiler.add_function(main)
        profiler.runcall(main)
        lstats = profiler.get_stats()

        stream = StringIO()
        stream.write('Timer unit: %g s\n\n' % lstats.unit)

        for (filename, lineno, name), timings in sorted(lstats.timings.iteritems()):
            if filename == '<string>':
                lineno = 2
                sublines = cell.splitlines()
            elif os.path.exists(filename) or filename.startswith('<ipython-input-'):
                if os.path.exists(filename):
                    linecache.clearcache()
                sublines = inspect.getblock(linecache.getlines(filename)[lineno-1:])
            else:
                raise IOError('failed to find {}'.format(filename))

            total_time = sum([time for _, _, time in timings])
            body = {}
            for l, nhits, time in timings:
                body[l-1] = (nhits, time, '%5.1f' % (time / nhits), '%5.1f' % (100 * time / total_time))
            template = '%6s %9s %12s %8s %8s  %-s'
            header = template % ('Line #', 'Hits', 'Time', 'Per Hit', '% Time', 'Line Contents')
            texts = []
            for l, line in enumerate(sublines, start=lineno):
                hits, time, per_hit, percent_time = body.get(l-1, ('', '', '', ''))
                texts.append(template % (l, hits, time, per_hit, percent_time, line.rstrip()))

            stream.write('Total time: %g s\n' % (total_time * lstats.unit))
            stream.write('File: %s\n' % filename)
            stream.write('Function: %s at line %s\n' % (name, lineno))
            stream.write('\n')
            stream.write(header)
            stream.write('\n')
            stream.write('=' * len(header))
            stream.write('\n')
            for text in texts:
                stream.write(text)
                stream.write('\n')
            stream.write('\n')

        page.page(stream.getvalue().rstrip())


def load_ipython_extension(ipython):
    ipython.register_magics(CPUProfiler)

Here's another version that's a bit less code (because it wraps the line magic, which is a hack). Thanks @auneri for the first cut!

• https://github.com/jdanbrown/potoo/blob/a7cc6a9/potoo/lprun_cell_magic.py