Random improvements

Author: pablogsal

Lib/profile/sample.py

@@ -4,6 +4,8 @@
 import time
 import _remote_debugging
 import argparse
+import _colorize
+from _colorize import ANSIColors
 
 
 class SampleProfile:
@@ -56,10 +58,148 @@ def sample(self, duration_sec=10):
 
         self.stats = self.convert_to_pstats(result)
 
-    def print_stats(self, sort=-1):
+    def get_color_for_value(self, value, max_value, theme):
+        """Return a color code based on the value's proportion of the maximum.
+        For profiling, higher values are more significant."""
+        if max_value == 0:
+            return ANSIColors.RESET
+        
+        ratio = value / max_value
+        if ratio > 0.8:
+            return ANSIColors.BOLD_RED  # Bright red for most significant values
+        elif ratio > 0.6:
+            return ANSIColors.RED  # Red for very significant values
+        elif ratio > 0.4:
+            return ANSIColors.YELLOW  # Yellow for significant values
+        elif ratio > 0.2:
+            return ANSIColors.GREEN  # Green for somewhat significant values
+        else:
+            return ANSIColors.RESET  # Default for less significant values
+
+    def print_stats(self, sort=-1, limit=None, show_summary=True):
         if not isinstance(sort, tuple):
             sort = (sort,)
-        pstats.SampledStats(self).strip_dirs().sort_stats(*sort).print_stats()
+        theme = _colorize.get_theme()
+        stats = pstats.SampledStats(self).strip_dirs()
+        
+        # Get the stats data
+        stats_list = []
+        for func, (cc, nc, tt, ct, callers) in stats.stats.items():
+            stats_list.append((func, cc, nc, tt, ct, callers))
+        
+        # Sort based on the requested field
+        sort_field = sort[0]
+        if sort_field == -1:  # stdname
+            stats_list.sort(key=lambda x: str(x[0]))
+        elif sort_field == 0:  # calls
+            stats_list.sort(key=lambda x: x[2], reverse=True)
+        elif sort_field == 1:  # time
+            stats_list.sort(key=lambda x: x[3], reverse=True)
+        elif sort_field == 2:  # cumulative
+            stats_list.sort(key=lambda x: x[4], reverse=True)
+        elif sort_field == 3:  # percall
+            stats_list.sort(key=lambda x: x[3]/x[2] if x[2] > 0 else 0, reverse=True)
+        elif sort_field == 4:  # cumpercall
+            stats_list.sort(key=lambda x: x[4]/x[2] if x[2] > 0 else 0, reverse=True)
+        
+        # Apply limit if specified
+        if limit is not None:
+            stats_list = stats_list[:limit]
+        
+        # Find the maximum values for each column to determine units
+        max_tt = max((tt for _, _, _, tt, _, _ in stats_list), default=0)
+        max_ct = max((ct for _, _, _, _, ct, _ in stats_list), default=0)
+        
+        # Determine appropriate units and format strings
+        if max_tt >= 1.0:
+            tt_unit = "s"
+            tt_scale = 1.0
+        elif max_tt >= 0.001:
+            tt_unit = "ms"
+            tt_scale = 1000.0
+        else:
+            tt_unit = "μs"
+            tt_scale = 1000000.0
+            
+        if max_ct >= 1.0:
+            ct_unit = "s"
+            ct_scale = 1.0
+        elif max_ct >= 0.001:
+            ct_unit = "ms"
+            ct_scale = 1000.0
+        else:
+            ct_unit = "μs"
+            ct_scale = 1000000.0
+        
+        # Print header with colors and units
+        header = (
+            f"{ANSIColors.BOLD_BLUE}Profile Stats:{ANSIColors.RESET}\n"
+            f"{ANSIColors.BOLD_BLUE}ncalls{ANSIColors.RESET}      "
+            f"{ANSIColors.BOLD_BLUE}tottime ({tt_unit}){ANSIColors.RESET}  "
+            f"{ANSIColors.BOLD_BLUE}percall ({tt_unit}){ANSIColors.RESET}  "
+            f"{ANSIColors.BOLD_BLUE}cumtime ({ct_unit}){ANSIColors.RESET}  "
+            f"{ANSIColors.BOLD_BLUE}percall ({ct_unit}){ANSIColors.RESET}  "
+            f"{ANSIColors.BOLD_BLUE}filename:lineno(function){ANSIColors.RESET}"
+        )
+        print(header)
+        
+        # Print each line with colors
+        for func, cc, nc, tt, ct, callers in stats_list:
+            if nc != cc:
+                ncalls = f"{nc}/{cc}"
+            else:
+                ncalls = str(nc)
+            
+            # Format numbers with proper alignment and precision (no colors)
+            tottime = f"{tt * tt_scale:8.3f}"
+            percall = f"{(tt/nc) * tt_scale:8.3f}" if nc > 0 else "    N/A"
+            cumtime = f"{ct * ct_scale:8.3f}"
+            cumpercall = f"{(ct/nc) * ct_scale:8.3f}" if nc > 0 else "    N/A"
+            
+            # Format the function name with colors
+            func_name = (
+                f"{ANSIColors.GREEN}{func[0]}{ANSIColors.RESET}:"
+                f"{ANSIColors.YELLOW}{func[1]}{ANSIColors.RESET}("
+                f"{ANSIColors.CYAN}{func[2]}{ANSIColors.RESET})"
+            )
+            
+            # Print the formatted line
+            print(
+                f"{ncalls:>8}  {tottime}  {percall}  "
+                f"{cumtime}  {cumpercall}  {func_name}"
+            )
+        
+        # Print summary of interesting functions if enabled
+        if show_summary and stats_list:
+            print(f"\n{ANSIColors.BOLD_BLUE}Summary of Interesting Functions:{ANSIColors.RESET}")
+            
+            # Most time-consuming functions (by total time)
+            print(f"\n{ANSIColors.BOLD_BLUE}Most Time-Consuming Functions:{ANSIColors.RESET}")
+            for func, _, nc, tt, ct, _ in sorted(stats_list, key=lambda x: x[3], reverse=True)[:3]:
+                if tt > 0:
+                    func_name = f"{func[0]}:{func[1]}({func[2]})"
+                    print(f"  {tt * tt_scale:8.3f}{tt_unit} total time, {(tt/nc) * tt_scale:8.3f}{tt_unit} per call: {func_name}")
+            
+            # Most called functions
+            print(f"\n{ANSIColors.BOLD_BLUE}Most Called Functions:{ANSIColors.RESET}")
+            for func, _, nc, tt, ct, _ in sorted(stats_list, key=lambda x: x[2], reverse=True)[:3]:
+                if nc > 0:
+                    func_name = f"{func[0]}:{func[1]}({func[2]})"
+                    print(f"  {nc:8d} calls, {(tt/nc) * tt_scale:8.3f}{tt_unit} per call: {func_name}")
+            
+            # Functions with highest per-call overhead
+            print(f"\n{ANSIColors.BOLD_BLUE}Functions with Highest Per-Call Overhead:{ANSIColors.RESET}")
+            for func, _, nc, tt, ct, _ in sorted(stats_list, key=lambda x: x[3]/x[2] if x[2] > 0 else 0, reverse=True)[:3]:
+                if nc > 0 and tt > 0:
+                    func_name = f"{func[0]}:{func[1]}({func[2]})"
+                    print(f"  {(tt/nc) * tt_scale:8.3f}{tt_unit} per call, {nc:8d} calls: {func_name}")
+            
+            # Functions with highest cumulative impact
+            print(f"\n{ANSIColors.BOLD_BLUE}Functions with Highest Cumulative Impact:{ANSIColors.RESET}")
+            for func, _, nc, tt, ct, _ in sorted(stats_list, key=lambda x: x[4], reverse=True)[:3]:
+                if ct > 0:
+                    func_name = f"{func[0]}:{func[1]}({func[2]})"
+                    print(f"  {ct * ct_scale:8.3f}{ct_unit} cumulative time, {(ct/nc) * ct_scale:8.3f}{ct_unit} per call: {func_name}")
 
     def dump_stats(self, file):
         stats_with_marker = dict(self.stats)
@@ -121,18 +261,31 @@ def sample(
     duration_sec=10,
     filename=None,
     all_threads=False,
+    limit=None,
+    show_summary=True,
 ):
     profile = SampleProfile(pid, sample_interval_usec, all_threads=False)
     profile.sample(duration_sec)
     if filename:
         profile.dump_stats(filename)
     else:
-        profile.print_stats(sort)
+        profile.print_stats(sort, limit, show_summary)
 
 
 def main():
     parser = argparse.ArgumentParser(
-        description="Sample a process's stack frames.", color=True
+        description=(
+            "Sample a process's stack frames.\n\n"
+            "Sort options:\n"
+            "  --sort-calls      Sort by number of calls (most called functions first)\n"
+            "  --sort-time       Sort by total time (most time-consuming functions first)\n"
+            "  --sort-cumulative Sort by cumulative time (functions with highest total impact first)\n"
+            "  --sort-percall    Sort by time per call (functions with highest per-call overhead first)\n"
+            "  --sort-cumpercall Sort by cumulative time per call (functions with highest cumulative overhead per call)\n"
+            "  --sort-name       Sort by function name (alphabetical order)\n\n"
+            "The default sort is by cumulative time (--sort-cumulative)."
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
     )
     parser.add_argument("pid", type=int, help="Process ID to sample.")
     parser.add_argument(
@@ -156,14 +309,86 @@ def main():
         help="Sample all threads in the process",
     )
     parser.add_argument("-o", "--outfile", help="Save stats to <outfile>")
+    parser.add_argument(
+        "--no-color",
+        action="store_true",
+        help="Disable color output",
+    )
+    parser.add_argument(
+        "-l",
+        "--limit",
+        type=int,
+        help="Limit the number of rows in the output",
+    )
+    parser.add_argument(
+        "--no-summary",
+        action="store_true",
+        help="Disable the summary section at the end of the output",
+    )
+    
+    # Add sorting options
+    sort_group = parser.add_mutually_exclusive_group()
+    sort_group.add_argument(
+        "--sort-calls",
+        action="store_const",
+        const=0,
+        dest="sort",
+        help="Sort by number of calls (most called functions first)",
+    )
+    sort_group.add_argument(
+        "--sort-time",
+        action="store_const",
+        const=1,
+        dest="sort",
+        help="Sort by total time (most time-consuming functions first)",
+    )
+    sort_group.add_argument(
+        "--sort-cumulative",
+        action="store_const",
+        const=2,
+        dest="sort",
+        help="Sort by cumulative time (functions with highest total impact first)",
+    )
+    sort_group.add_argument(
+        "--sort-percall",
+        action="store_const",
+        const=3,
+        dest="sort",
+        help="Sort by time per call (functions with highest per-call overhead first)",
+    )
+    sort_group.add_argument(
+        "--sort-cumpercall",
+        action="store_const",
+        const=4,
+        dest="sort",
+        help="Sort by cumulative time per call (functions with highest cumulative overhead per call)",
+    )
+    sort_group.add_argument(
+        "--sort-name",
+        action="store_const",
+        const=5,
+        dest="sort",
+        help="Sort by function name (alphabetical order)",
+    )
+    
+    # Set default sort to cumulative time
+    parser.set_defaults(sort=2)
+    
     args = parser.parse_args()
 
+    # Set color theme based on --no-color flag
+    if args.no_color:
+        _colorize.set_theme(_colorize.theme_no_color)
+
     sample(
         args.pid,
         sample_interval_usec=args.interval,
         duration_sec=args.duration,
         filename=args.outfile,
         all_threads=args.all_threads,
+        limit=args.limit,
+        sort=args.sort,
+        show_summary=not args.no_summary,
     )