Improve CLI

Author: pablogsal

Lib/profile/sample.py

@@ -166,70 +166,94 @@ def _print_top_functions(stats_list, title, key_func, format_line, n=3):
             f"\n{ANSIColors.BOLD_BLUE}Summary of Interesting Functions:{ANSIColors.RESET}"
         )
 
+        # Aggregate stats by fully qualified function name (ignoring line numbers)
+        func_aggregated = {}
+        for func, prim_calls, total_calls, total_time, cumulative_time, callers in stats_list:
+            # Use filename:function_name as the key to get fully qualified name
+            qualified_name = f"{func[0]}:{func[2]}"
+            if qualified_name not in func_aggregated:
+                func_aggregated[qualified_name] = [0, 0, 0, 0]  # prim_calls, total_calls, total_time, cumulative_time
+            func_aggregated[qualified_name][0] += prim_calls
+            func_aggregated[qualified_name][1] += total_calls
+            func_aggregated[qualified_name][2] += total_time
+            func_aggregated[qualified_name][3] += cumulative_time
+
+        # Convert aggregated data back to list format for processing
+        aggregated_stats = []
+        for qualified_name, (prim_calls, total_calls, total_time, cumulative_time) in func_aggregated.items():
+            # Parse the qualified name back to filename and function name
+            if ":" in qualified_name:
+                filename, func_name = qualified_name.rsplit(":", 1)
+            else:
+                filename, func_name = "", qualified_name
+            # Create a dummy func tuple with filename and function name for display
+            dummy_func = (filename, "", func_name)
+            aggregated_stats.append((dummy_func, prim_calls, total_calls, total_time, cumulative_time, {}))
+
         # Most time-consuming functions (by total time)
         def format_time_consuming(stat):
-            func, _, nc, tt, _, _ = stat
-            if tt > 0:
+            func, _, total_calls, total_time, _, _ = stat
+            if total_time > 0:
                 return (
-                    f"{tt * tt_scale:8.3f} {tt_unit} total time, "
-                    f"{(tt / nc) * tt_scale:8.3f} {tt_unit} per call: {_format_func_name(func)}"
+                    f"{total_time * tt_scale:8.3f} {tt_unit} total time, "
+                    f"{(total_time / total_calls) * tt_scale:8.3f} {tt_unit} per call: {_format_func_name(func)}"
                 )
             return None
 
         _print_top_functions(
-            stats_list,
+            aggregated_stats,
             "Most Time-Consuming Functions",
             key_func=lambda x: x[3],
             format_line=format_time_consuming,
         )
 
         # Most called functions
         def format_most_called(stat):
-            func, _, nc, tt, _, _ = stat
-            if nc > 0:
+            func, _, total_calls, total_time, _, _ = stat
+            if total_calls > 0:
                 return (
-                    f"{nc:8d} calls, {(tt / nc) * tt_scale:8.3f} {tt_unit} "
+                    f"{total_calls:8d} calls, {(total_time / total_calls) * tt_scale:8.3f} {tt_unit} "
                     f"per call: {_format_func_name(func)}"
                 )
             return None
 
         _print_top_functions(
-            stats_list,
+            aggregated_stats,
             "Most Called Functions",
             key_func=lambda x: x[2],
             format_line=format_most_called,
         )
 
         # Functions with highest per-call overhead
         def format_overhead(stat):
-            func, _, nc, tt, _, _ = stat
-            if nc > 0 and tt > 0:
+            func, _, total_calls, total_time, _, _ = stat
+            if total_calls > 0 and total_time > 0:
                 return (
-                    f"{(tt / nc) * tt_scale:8.3f} {tt_unit} per call, "
-                    f"{nc:8d} calls: {_format_func_name(func)}"
+                    f"{(total_time / total_calls) * tt_scale:8.3f} {tt_unit} per call, "
+                    f"{total_calls:8d} calls: {_format_func_name(func)}"
                 )
             return None
 
         _print_top_functions(
-            stats_list,
+            aggregated_stats,
             "Functions with Highest Per-Call Overhead",
             key_func=lambda x: x[3] / x[2] if x[2] > 0 else 0,
             format_line=format_overhead,
         )
 
         # Functions with highest cumulative impact
         def format_cumulative(stat):
-            func, _, nc, _, ct, _ = stat
-            if ct > 0:
+            func, _, total_calls, _, cumulative_time, _ = stat
+            if cumulative_time > 0:
                 return (
-                    f"{ct * ct_scale:8.3f} {ct_unit} cumulative time, "
-                    f"{(ct / nc) * ct_scale:8.3f} {ct_unit} per call: "
+                    f"{cumulative_time * ct_scale:8.3f} {ct_unit} cumulative time, "
+                    f"{(cumulative_time / total_calls) * ct_scale:8.3f} {ct_unit} per call: "
                     f"{_format_func_name(func)}"
                 )
             return None
 
         _print_top_functions(
-            stats_list,
+            aggregated_stats,
             "Functions with Highest Cumulative Impact",
             key_func=lambda x: x[4],
             format_line=format_cumulative,
@@ -270,125 +294,185 @@ def sample(
     else:
         collector.export(filename)
 
+def _validate_collapsed_format_args(args, parser):
+    # Check for incompatible pstats options
+    invalid_opts = []
+    
+    # Get list of pstats-specific options
+    pstats_options = {
+        'sort': None,
+        'limit': None,
+        'no_summary': False
+    }
+
+    # Find the default values from the argument definitions
+    for action in parser._actions:
+        if action.dest in pstats_options and hasattr(action, 'default'):
+            pstats_options[action.dest] = action.default
+
+    # Check if any pstats-specific options were provided by comparing with defaults
+    for opt, default in pstats_options.items():
+        if getattr(args, opt) != default:
+            invalid_opts.append(opt.replace('no_', ''))
+    
+    if invalid_opts:
+        parser.error(f"The following options are only valid with --pstats format: {', '.join(invalid_opts)}")
+    
+    # Set default output filename for collapsed format
+    if not args.outfile:
+        args.outfile = f"collapsed.{args.pid}.txt"
+
 
 def main():
+    # Create the main parser
     parser = argparse.ArgumentParser(
         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)."
-            "Format descriptions:\n"
-            "  pstats     Standard Python profiler output format\n"
-            "  collapsed  Stack traces in collapsed format (file:function:line;file:function:line;... count)\n"
-            "             Useful for generating flamegraphs with tools like flamegraph.pl"
+            "Sample a process's stack frames and generate profiling data.\n"
+            "Supports two output formats:\n"
+            "  - pstats: Detailed profiling statistics with sorting options\n" 
+            "  - collapsed: Stack traces for generating flamegraphs\n"
+            "\n"
+            "Examples:\n"
+            "  # Profile process 1234 for 10 seconds with default settings\n"
+            "  python -m profile.sample 1234\n"
+            "\n"
+            "  # Profile with custom interval and duration, save to file\n" 
+            "  python -m profile.sample -i 50 -d 30 -o profile.stats 1234\n"
+            "\n"
+            "  # Generate collapsed stacks for flamegraph\n"
+            "  python -m profile.sample --collapsed 1234\n"
+            "\n"
+            "  # Profile all threads, sort by total time\n"
+            "  python -m profile.sample -a --sort-time 1234\n"
+            "\n"
+            "  # Profile for 1 minute with 1ms sampling interval\n"
+            "  python -m profile.sample -i 1000 -d 60 1234\n"
+            "\n"
+            "  # Show only top 20 functions sorted by calls\n"
+            "  python -m profile.sample --sort-calls -l 20 1234\n"
+            "\n"
+            "  # Profile all threads and save collapsed stacks\n"
+            "  python -m profile.sample -a --collapsed -o stacks.txt 1234"
         ),
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        color=True,
+        formatter_class=argparse.RawDescriptionHelpFormatter
     )
-    parser.add_argument("pid", type=int, help="Process ID to sample.")
+
+    # Required arguments
     parser.add_argument(
-        "-i",
-        "--interval",
+        "pid", 
         type=int,
-        default=10,
-        help="Sampling interval in microseconds (default: 10 usec)",
+        help="Process ID to sample"
     )
-    parser.add_argument(
-        "-d",
-        "--duration",
+
+    # Sampling options
+    sampling_group = parser.add_argument_group("Sampling configuration")
+    sampling_group.add_argument(
+        "-i", "--interval",
         type=int,
-        default=10,
-        help="Sampling duration in seconds (default: 10 seconds)",
-    )
-    parser.add_argument(
-        "-a",
-        "--all-threads",
-        action="store_true",
-        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",
+        default=100,
+        help="Sampling interval in microseconds (default: 100)"
     )
-    parser.add_argument(
-        "-l",
-        "--limit",
+    sampling_group.add_argument(
+        "-d", "--duration",
         type=int,
-        help="Limit the number of rows in the output",
+        default=10,
+        help="Sampling duration in seconds (default: 10)"
     )
-    parser.add_argument(
-        "--no-summary",
+    sampling_group.add_argument(
+        "-a", "--all-threads",
         action="store_true",
-        help="Disable the summary section at the end of the output",
+        help="Sample all threads in the process instead of just the main thread"
     )
-    parser.add_argument(
-        "--format",
-        choices=["pstats", "collapsed"],
+
+    # Output format selection
+    output_group = parser.add_argument_group("Output options")
+    output_format = output_group.add_mutually_exclusive_group()
+    output_format.add_argument(
+        "--pstats",
+        action="store_const",
+        const="pstats",
+        dest="format",
         default="pstats",
-        help="Output format (default: pstats)",
+        help="Generate pstats output (default)"
+    )
+    output_format.add_argument(
+        "--collapsed",
+        action="store_const", 
+        const="collapsed",
+        dest="format",
+        help="Generate collapsed stack traces for flamegraphs"
+    )
+    
+    output_group.add_argument(
+        "-o", "--outfile",
+        help="Save output to a file (if omitted, prints to stdout for pstats, "
+             "or saves to collapsed.<pid>.txt for collapsed format)"
     )
 
-    # Add sorting options
-    sort_group = parser.add_mutually_exclusive_group()
+    # pstats-specific options
+    pstats_group = parser.add_argument_group("pstats format options")
+    sort_group = pstats_group.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)",
+        help="Sort by number of calls"
     )
     sort_group.add_argument(
         "--sort-time",
         action="store_const",
         const=1,
         dest="sort",
-        help="Sort by total time (most time-consuming functions first)",
+        help="Sort by total time"
     )
     sort_group.add_argument(
         "--sort-cumulative",
         action="store_const",
         const=2,
         dest="sort",
-        help="Sort by cumulative time (functions with highest total impact first)",
+        default=2,
+        help="Sort by cumulative time (default)"
     )
     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)",
+        help="Sort by time per call"
     )
     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)",
+        help="Sort by cumulative time per call"
     )
     sort_group.add_argument(
         "--sort-name",
         action="store_const",
         const=5,
         dest="sort",
-        help="Sort by function name (alphabetical order)",
+        help="Sort by function name"
     )
 
-    # Set default sort to cumulative time
-    parser.set_defaults(sort=2)
+    pstats_group.add_argument(
+        "-l", "--limit",
+        type=int,
+        help="Limit the number of rows in the output",
+        default=15,
+    )
+    pstats_group.add_argument(
+        "--no-summary",
+        action="store_true",
+        help="Disable the summary section in the output"
+    )
 
     args = parser.parse_args()
 
-    # Set color theme based on --no-color flag
-    if args.no_color:
-        _colorize.set_theme(_colorize.theme_no_color)
+    # Validate format-specific arguments
+    if args.format == "collapsed":
+        _validate_collapsed_format_args(args, parser)
 
     sample(
         args.pid,
@@ -401,7 +485,5 @@ def main():
         show_summary=not args.no_summary,
         output_format=args.format,
     )
-
-
 if __name__ == "__main__":
     main()

Lib/test/test_sample_profiler.py

@@ -535,12 +535,12 @@ def test_argument_parsing_basic(self):
 
             mock_sample.assert_called_once_with(
                 12345,
-                sample_interval_usec=10,
+                sample_interval_usec=100,
                 duration_sec=10,
                 filename=None,
                 all_threads=False,
-                limit=None,
-                sort=2,
+                limit=15,
+                sort=None,
                 show_summary=True,
                 output_format="pstats",
             )