databricks line_profiler

databricks_line-profiler.py

dbutils.library.installPyPI("line_profiler") # this fails with Python 3.7

import inspect
from io import StringIO
from line_profiler import LineProfiler

def profile_function(my_func, args, *kwargs):
    lp = LineProfiler()
    output_val = lp(my_func)(args, *kwargs)
    mystdout = StringIO()
    lp.print_stats(stream=mystdout)
    #Redirect stdout so we can grab profile output
    lprof_lines = mystdout.getvalue().split('\n')
    profile_start = 1 + next(idx for idx, line in enumerate(lprof_lines) if '=====' in line)
    lprof_code_lines = lprof_lines[profile_start:-1]

    source_lines = inspect.getsource(my_func).split('\n')

    if len(source_lines) != len(lprof_code_lines):
        print("WARNING! Mismatch in source length and returned line profiler estimates")
        print('\n'.join(lprof_lines))
        print("---- Code ----")
        print(source)
    else:
        print("\n".join(lprof_lines[:profile_start]))
        print("\n".join(["{0} \t {1}".format(l, s) for l, s in zip(lprof_code_lines, source_lines)]))

    return output_val