Filter the stats for our modules only and sort by average time per function

filter_and_sort_by_avg.py

def is_my_key(key, modules):
  """ Returns True if the key belongs to one of the modules.

  Args:
    key: key of the stats dictionary.
    modules: list of strings with the modules to be matched.
  """
  fn = key[0]
  for m in modules:
    module_name = '%s' % m
    if module_name in fn:
      return True
  return False


stats_dict = p.sort_stats('cumulative').stats

mymodules = ['dofns', 'pipeline']
mykeys = [k for k in stats_dict.keys() if is_my_key(k, mymodules)]
mystats = {k: stats_dict[k] for k in mykeys}

percall = {}

for k,v in mystats.items():
  _, ncalls, _, cumtime, _ = v
  avg_time = cumtime/ncalls
  percall[k] = avg_time

sorted(percall.items(), key=lambda x: x[1], reverse=True)