CLI update

Author: pablogsal

Lib/profiling/sampling/cli.py

@@ -197,11 +197,8 @@ def _add_sampling_options(parser):
     )
     sampling_group.add_argument(
         "--async-aware",
-        choices=["running", "all"],
-        default=None,
-        metavar="MODE",
-        help='Enable async-aware profiling: "running" (only running task) '
-        'or "all" (all tasks including waiting)',
+        action="store_true",
+        help="Enable async-aware profiling (uses task-based stack reconstruction)",
     )
 
 
@@ -213,7 +210,14 @@ def _add_mode_options(parser):
         choices=["wall", "cpu", "gil"],
         default="wall",
         help="Sampling mode: wall (all samples), cpu (only samples when thread is on CPU), "
-        "gil (only samples when thread holds the GIL)",
+        "gil (only samples when thread holds the GIL). Incompatible with --async-aware",
+    )
+    mode_group.add_argument(
+        "--async-mode",
+        choices=["running", "all"],
+        default="running",
+        help='Async profiling mode: "running" (only running task) '
+        'or "all" (all tasks including waiting). Requires --async-aware',
     )
 
 
@@ -391,7 +395,7 @@ def _validate_args(args, parser):
         )
 
     # Async-aware mode is incompatible with --native, --no-gc, --mode, and --all-threads
-    if args.async_aware is not None:
+    if args.async_aware:
         issues = []
         if args.native:
             issues.append("--native")
@@ -407,6 +411,10 @@ def _validate_args(args, parser):
                 "Async-aware profiling uses task-based stack reconstruction."
             )
 
+    # --async-mode requires --async-aware
+    if hasattr(args, 'async_mode') and args.async_mode != "running" and not args.async_aware:
+        parser.error("--async-mode requires --async-aware to be enabled.")
+
     # Live mode is incompatible with format options
     if hasattr(args, 'live') and args.live:
         if args.format != "pstats":
@@ -595,7 +603,7 @@ def _handle_attach(args):
         all_threads=args.all_threads,
         realtime_stats=args.realtime_stats,
         mode=mode,
-        async_aware=args.async_aware,
+        async_aware=args.async_mode if args.async_aware else None,
         native=args.native,
         gc=args.gc,
     )
@@ -644,7 +652,7 @@ def _handle_run(args):
             all_threads=args.all_threads,
             realtime_stats=args.realtime_stats,
             mode=mode,
-            async_aware=args.async_aware,
+            async_aware=args.async_mode if args.async_aware else None,
             native=args.native,
             gc=args.gc,
         )
@@ -677,7 +685,7 @@ def _handle_live_attach(args, pid):
         limit=20,  # Default limit
         pid=pid,
         mode=mode,
-        async_aware=args.async_aware,
+        async_aware=args.async_mode if args.async_aware else None,
     )
 
     # Sample in live mode
@@ -688,7 +696,7 @@ def _handle_live_attach(args, pid):
         all_threads=args.all_threads,
         realtime_stats=args.realtime_stats,
         mode=mode,
-        async_aware=args.async_aware,
+        async_aware=args.async_mode if args.async_aware else None,
         native=args.native,
         gc=args.gc,
     )
@@ -718,7 +726,7 @@ def _handle_live_run(args):
         limit=20,  # Default limit
         pid=process.pid,
         mode=mode,
-        async_aware=args.async_aware,
+        async_aware=args.async_mode if args.async_aware else None,
     )
 
     # Profile the subprocess in live mode
@@ -730,7 +738,7 @@ def _handle_live_run(args):
             all_threads=args.all_threads,
             realtime_stats=args.realtime_stats,
             mode=mode,
-            async_aware=args.async_aware,
+            async_aware=args.async_mode if args.async_aware else None,
             native=args.native,
             gc=args.gc,
         )

Lib/test/test_profiling/test_sampling_profiler/test_cli.py

@@ -548,9 +548,9 @@ def test_sort_options(self):
                 mock_sample.assert_called_once()
                 mock_sample.reset_mock()
 
-    def test_async_aware_argument_all(self):
-        """Test --async-aware all argument is parsed correctly."""
-        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "all"]
+    def test_async_aware_flag_defaults_to_running(self):
+        """Test --async-aware flag enables async profiling with default 'running' mode."""
+        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware"]
 
         with (
             mock.patch("sys.argv", test_args),
@@ -560,13 +560,13 @@ def test_async_aware_argument_all(self):
             main()
 
             mock_sample.assert_called_once()
-            # Verify async_aware was passed
+            # Verify async_aware was passed with default "running" mode
             call_kwargs = mock_sample.call_args[1]
-            self.assertEqual(call_kwargs.get("async_aware"), "all")
+            self.assertEqual(call_kwargs.get("async_aware"), "running")
 
-    def test_async_aware_argument_running(self):
-        """Test --async-aware running argument is parsed correctly."""
-        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "running"]
+    def test_async_aware_with_async_mode_all(self):
+        """Test --async-aware with --async-mode all."""
+        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--async-mode", "all"]
 
         with (
             mock.patch("sys.argv", test_args),
@@ -577,10 +577,10 @@ def test_async_aware_argument_running(self):
 
             mock_sample.assert_called_once()
             call_kwargs = mock_sample.call_args[1]
-            self.assertEqual(call_kwargs.get("async_aware"), "running")
+            self.assertEqual(call_kwargs.get("async_aware"), "all")
 
     def test_async_aware_default_is_none(self):
-        """Test --async-aware defaults to None when not specified."""
+        """Test async_aware defaults to None when --async-aware not specified."""
         test_args = ["profiling.sampling.cli", "attach", "12345"]
 
         with (
@@ -594,9 +594,9 @@ def test_async_aware_default_is_none(self):
             call_kwargs = mock_sample.call_args[1]
             self.assertIsNone(call_kwargs.get("async_aware"))
 
-    def test_async_aware_invalid_choice(self):
-        """Test --async-aware with invalid choice raises error."""
-        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "invalid"]
+    def test_async_mode_invalid_choice(self):
+        """Test --async-mode with invalid choice raises error."""
+        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--async-mode", "invalid"]
 
         with (
             mock.patch("sys.argv", test_args),
@@ -608,9 +608,25 @@ def test_async_aware_invalid_choice(self):
 
         self.assertEqual(cm.exception.code, 2)  # argparse error
 
+    def test_async_mode_requires_async_aware(self):
+        """Test --async-mode without --async-aware raises error."""
+        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-mode", "all"]
+
+        with (
+            mock.patch("sys.argv", test_args),
+            mock.patch("sys.stderr", io.StringIO()) as mock_stderr,
+            self.assertRaises(SystemExit) as cm,
+        ):
+            from profiling.sampling.cli import main
+            main()
+
+        self.assertEqual(cm.exception.code, 2)  # argparse error
+        error_msg = mock_stderr.getvalue()
+        self.assertIn("--async-mode requires --async-aware", error_msg)
+
     def test_async_aware_incompatible_with_native(self):
         """Test --async-aware is incompatible with --native."""
-        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "all", "--native"]
+        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--native"]
 
         with (
             mock.patch("sys.argv", test_args),
@@ -627,7 +643,7 @@ def test_async_aware_incompatible_with_native(self):
 
     def test_async_aware_incompatible_with_no_gc(self):
         """Test --async-aware is incompatible with --no-gc."""
-        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "running", "--no-gc"]
+        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--no-gc"]
 
         with (
             mock.patch("sys.argv", test_args),
@@ -644,7 +660,7 @@ def test_async_aware_incompatible_with_no_gc(self):
 
     def test_async_aware_incompatible_with_both_native_and_no_gc(self):
         """Test --async-aware is incompatible with both --native and --no-gc."""
-        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "all", "--native", "--no-gc"]
+        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--native", "--no-gc"]
 
         with (
             mock.patch("sys.argv", test_args),
@@ -662,7 +678,7 @@ def test_async_aware_incompatible_with_both_native_and_no_gc(self):
 
     def test_async_aware_incompatible_with_mode(self):
         """Test --async-aware is incompatible with --mode (non-wall)."""
-        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "all", "--mode", "cpu"]
+        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--mode", "cpu"]
 
         with (
             mock.patch("sys.argv", test_args),
@@ -679,7 +695,7 @@ def test_async_aware_incompatible_with_mode(self):
 
     def test_async_aware_incompatible_with_all_threads(self):
         """Test --async-aware is incompatible with --all-threads."""
-        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "running", "--all-threads"]
+        test_args = ["profiling.sampling.cli", "attach", "12345", "--async-aware", "--all-threads"]
 
         with (
             mock.patch("sys.argv", test_args),