File size: 2,232 Bytes
079c32c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import atexit
import pstats
import io
import cProfile
import os
def register_profiler(write_profile, pr, folder_path):
atexit.register(write_profile, pr, folder_path)
class Profiler:
"""
Overview:
A class for profiling code execution. It can be used as a context manager or a decorator.
Interfaces:
``__init__``, ``mkdir``, ``write_profile``, ``profile``.
"""
def __init__(self):
"""
Overview:
Initialize the Profiler object.
"""
self.pr = cProfile.Profile()
def mkdir(self, directory: str):
"""
OverView:
Create a directory if it doesn't exist.
Arguments:
- directory (:obj:`str`): The path of the directory to be created.
"""
if not os.path.exists(directory):
os.makedirs(directory)
def write_profile(self, pr: cProfile.Profile, folder_path: str):
"""
OverView:
Write the profiling results to files.
Arguments:
- pr (:obj:`cProfile.Profile`): The profiler object containing the profiling results.
- folder_path (:obj:`str`): The path of the folder where the profiling files will be saved.
"""
pr.disable()
s_tottime = io.StringIO()
s_cumtime = io.StringIO()
ps = pstats.Stats(pr, stream=s_tottime).sort_stats('tottime')
ps.print_stats()
with open(folder_path + "/profile_tottime.txt", 'w+') as f:
f.write(s_tottime.getvalue())
ps = pstats.Stats(pr, stream=s_cumtime).sort_stats('cumtime')
ps.print_stats()
with open(folder_path + "/profile_cumtime.txt", 'w+') as f:
f.write(s_cumtime.getvalue())
pr.dump_stats(folder_path + "/profile.prof")
def profile(self, folder_path="./tmp"):
"""
OverView:
Enable profiling and save the results to files.
Arguments:
- folder_path (:obj:`str`): The path of the folder where the profiling files will be saved. \
Defaults to "./tmp".
"""
self.mkdir(folder_path)
self.pr.enable()
register_profiler(self.write_profile, self.pr, folder_path)
|