Add features to profiling utilities

PyMca5/PyMcaMisc/ProfilingUtils.py

@@ -1,4 +1,4 @@
-#/*##########################################################################
+# /*##########################################################################
 #
 # The PyMca X-Ray Fluorescence Toolkit
 #
@@ -30,63 +30,102 @@
 __contact__ = "[email protected]"
 __license__ = "MIT"
 __copyright__ = "European Synchrotron Radiation Facility, Grenoble, France"
+
 try:
     import tracemalloc
     import linecache
 except ImportError:
     tracemalloc = None
 import cProfile
 import pstats
+
 try:
     from StringIO import StringIO
 except ImportError:
     from io import StringIO
-import os
 import logging
 from contextlib import contextmanager
 
 
 logger = logging.getLogger(__name__)
 
 
-def print_malloc_snapshot(snapshot, key_type='lineno', limit=10, units='KB'):
+class bcolors:
+    HEADER = "\033[95m"
+    OKBLUE = "\033[94m"
+    OKGREEN = "\033[92m"
+    WARNING = "\033[93m"
+    FAIL = "\033[91m"
+    ENDC = "\033[0m"
+    BOLD = "\033[1m"
+    UNDERLINE = "\033[4m"
+
+
+def durationfmtcolor(x):
+    if x < 0.0005:
+        return bcolors.OKGREEN + ("%6dµs" % (x * 1000000)) + bcolors.ENDC
+    elif x < 0.001:
+        return bcolors.WARNING + ("%6dµs" % (x * 1000000)) + bcolors.ENDC
+    elif x < 0.1:
+        return bcolors.WARNING + ("%6dms" % (x * 1000)) + bcolors.ENDC
+    else:
+        return bcolors.FAIL + ("%8.3f" % x) + bcolors.ENDC
+
+
+def durationfmt(x):
+    if x < 0.001:
+        return "%6dµs" % (x * 1000000)
+    elif x < 0.1:
+        return "%6dms" % (x * 1000)
+    else:
+        return "%8.3f" % x
+
+
+def print_malloc_snapshot(snapshot, key_type="lineno", limit=10, units="KB"):
     """
     :param tracemalloc.Snapshot snapshot:
     :param str key_type:
     :param int limit: limit number of lines
     :param str units: B, KB, MB, GB
     """
-    n = ['b', 'kb', 'mb', 'gb'].index(units.lower())
-    sunits, units = units, 1024**n
+    n = ["b", "kb", "mb", "gb"].index(units.lower())
+    sunits, units = units, 1024 ** n
 
-    snapshot = snapshot.filter_traces((
-        tracemalloc.Filter(False, "<frozen importlib._bootstrap>"),
-        tracemalloc.Filter(False, "<unknown>"),
-    ))
+    snapshot = snapshot.filter_traces(
+        (
+            tracemalloc.Filter(False, "<frozen importlib._bootstrap>"),
+            tracemalloc.Filter(False, "<unknown>"),
+        )
+    )
     top_stats = snapshot.statistics(key_type)
     total = sum(stat.size for stat in top_stats)
 
-    print('================Memory profile================')
+    logger.info("================Memory profile================")
+    out = ""
     for index, stat in enumerate(top_stats, 1):
         frame = stat.traceback[0]
         # replace "/path/to/module/file.py" with "module/file.py"
-        #filename = os.sep.join(frame.filename.split(os.sep)[-2:])
+        # filename = os.sep.join(frame.filename.split(os.sep)[-2:])
         filename = frame.filename
-        print("#%s: %s:%s: %.1f %s"
-              % (index, filename, frame.lineno, stat.size / units, sunits))
+        out += "\n#%s: %s:%s: %.1f %s" % (
+            index,
+            filename,
+            frame.lineno,
+            stat.size / units,
+            sunits,
+        )
         line = linecache.getline(frame.filename, frame.lineno).strip()
         if line:
-            print('    %s' % line)
+            out += "\n    %s" % line
         if index >= limit:
             break
-
     other = top_stats[index:]
     if other:
         size = sum(stat.size for stat in other)
-        print("%s other: %.1f %s" % (len(other), size / units, sunits))
-
-    print("Total allocated size: %.1f %s" % (total / units, sunits))
-    print('============================================')
+        out += "\n%s other: %.1f %s" % (len(other), size / units, sunits)
+    out += "\nTotal allocated size: %.1f %s" % (total / units, sunits)
+    logger.info(out)
+    logger.info("============================================")
 
 
 @contextmanager
@@ -95,42 +134,89 @@ def print_malloc_context(**kwargs):
     :param **kwargs: see print_malloc_snapshot
     """
     if tracemalloc is None:
-        logger.error('tracemalloc required')
+        logger.info("tracemalloc required")
         return
     tracemalloc.start()
-    yield
-    snapshot = tracemalloc.take_snapshot()
-    print_malloc_snapshot(snapshot, **kwargs)
+    try:
+        yield
+    finally:
+        snapshot = tracemalloc.take_snapshot()
+        print_malloc_snapshot(snapshot, **kwargs)
 
 
 @contextmanager
-def print_time_context(restrictions=None, sortby='cumtime'):
+def print_time_context(timelimit=None, sortby="cumtime", color=False, filename=None):
+    """
+    :param int or float timelimit: number of lines or fraction (float between 0 and 1)
+    :param str sortby: sort time profile
+    :param bool color:
+    :param str filename:
+    """
     pr = cProfile.Profile()
     pr.enable()
-    yield
-    pr.disable()
-    s = StringIO()
-    ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
-    if restrictions is None:
-        restrictions = (0.1,)
-    ps.print_stats(*restrictions)
-    print('================Time profile================')
-    print(s.getvalue())
-    print('============================================')
+    try:
+        yield
+    finally:
+        pr.disable()
+        if isinstance(sortby, str):
+            sortby = [sortby]
+        if color:
+            pstats.f8 = durationfmtcolor
+        else:
+            pstats.f8 = durationfmt
+        for i, sortmethod in enumerate(sortby):
+            s = StringIO()
+            ps = pstats.Stats(pr, stream=s)
+            if sortmethod:
+                ps = ps.sort_stats(sortmethod)
+            if timelimit is None:
+                timelimit = (0.1,)
+            elif not isinstance(timelimit, tuple):
+                timelimit = (timelimit,)
+            ps.print_stats(*timelimit)
+            if filename and i == 0:
+                ps.dump_stats(filename)
+            logger.info("================Time profile================")
+            msg = "\n" + s.getvalue()
+            msg += "\n Saved as {}".format(repr(filename))
+            logger.info(msg)
+            logger.info("============================================")
 
 
 @contextmanager
-def profile(memory=True, time=True, memlimit=10,
-            restrictions=None, sortby='cumtime'):
+def profile(
+    memory=True,
+    time=True,
+    memlimit=10,
+    timelimit=None,
+    sortby="cumtime",
+    color=False,
+    filename=None,
+    units="KB",
+):
+    """
+    :param bool memory: profile memory usage
+    :param bool time: execution time
+    :param int memlimit: number of lines
+    :param int or float timelimit: number of lines or fraction (float between 0 and 1)
+    :param str sortby: sort time profile
+    :param bool color:
+    :param str filename: dump for visual tools
+    :param str units: memory units
+    """
     if not memory and not time:
-        yield
+        return
     elif memory and time:
-        with print_time_context(restrictions=restrictions, sortby=sortby):
-            with print_malloc_context(limit=memlimit):
+        with print_time_context(
+            timelimit=timelimit, sortby=sortby, color=color, filename=filename
+        ):
+            with print_malloc_context(limit=memlimit, units=units):
                 yield
     elif memory:
-        with print_malloc_context(limit=memlimit):
+        with print_malloc_context(limit=memlimit, units=units):
             yield
     else:
-        with print_time_context(restrictions=restrictions, sortby=sortby):
+        with print_time_context(
+            timelimit=timelimit, sortby=sortby, color=color, filename=filename
+        ):
             yield