def profile_ctx(engine='pyinstrument'): """ A context manager which profiles the body of the with statement with the supplied profiling engine and returns the profiling object in a list. Arguments --------- engine: str The profiling engine, e.g. 'pyinstrument' or 'snakeviz' Returns ------- sessions: list A list containing the profiling session. """ if engine == 'pyinstrument': from pyinstrument import Profiler prof = Profiler() prof.start() elif engine == 'snakeviz': prof = Profile() prof.enable() elif engine is None: pass sessions = [] yield sessions if engine == 'pyinstrument': sessions.append(prof.stop()) elif engine == 'snakeviz': prof.disable() sessions.append(prof)