diff --git a/Lib/profile/sample.py b/Lib/profile/sample.py index 97d23611e67ad7..0ef50c217cbf95 100644 --- a/Lib/profile/sample.py +++ b/Lib/profile/sample.py @@ -2,6 +2,7 @@ import _remote_debugging import os import pstats +import subprocess import statistics import sys import sysconfig @@ -542,46 +543,66 @@ def main(): parser = argparse.ArgumentParser( description=( "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" + "Supports the following target modes:\n" + " - -p PID: Profile an existing process by PID\n" + " - -m MODULE [ARGS...]: Profile a module as python -m module ... \n" + " - filename [ARGS...]: Profile the specified script by running it in a subprocess\n" "\n" "Examples:\n" " # Profile process 1234 for 10 seconds with default settings\n" - " python -m profile.sample 1234\n" + " python -m profile.sample -p 1234\n" + "\n" + " # Profile a script by running it in a subprocess\n" + " python -m profile.sample myscript.py arg1 arg2\n" + "\n" + " # Profile a module by running it as python -m module in a subprocess\n" + " python -m profile.sample -m mymodule arg1 arg2\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" + " python -m profile.sample -i 50 -d 30 -o profile.stats -p 1234\n" "\n" " # Generate collapsed stacks for flamegraph\n" - " python -m profile.sample --collapsed 1234\n" + " python -m profile.sample --collapsed -p 1234\n" "\n" " # Profile all threads, sort by total time\n" - " python -m profile.sample -a --sort-tottime 1234\n" + " python -m profile.sample -a --sort-tottime -p 1234\n" "\n" " # Profile for 1 minute with 1ms sampling interval\n" - " python -m profile.sample -i 1000 -d 60 1234\n" + " python -m profile.sample -i 1000 -d 60 -p 1234\n" "\n" " # Show only top 20 functions sorted by direct samples\n" - " python -m profile.sample --sort-nsamples -l 20 1234\n" + " python -m profile.sample --sort-nsamples -l 20 -p 1234\n" "\n" " # Profile all threads and save collapsed stacks\n" - " python -m profile.sample -a --collapsed -o stacks.txt 1234\n" + " python -m profile.sample -a --collapsed -o stacks.txt -p 1234\n" "\n" " # Profile with real-time sampling statistics\n" - " python -m profile.sample --realtime-stats 1234\n" + " python -m profile.sample --realtime-stats -p 1234\n" "\n" " # Sort by sample percentage to find most sampled functions\n" - " python -m profile.sample --sort-sample-pct 1234\n" + " python -m profile.sample --sort-sample-pct -p 1234\n" "\n" " # Sort by cumulative samples to find functions most on call stack\n" - " python -m profile.sample --sort-nsamples-cumul 1234" + " python -m profile.sample --sort-nsamples-cumul -p 1234\n" ), formatter_class=argparse.RawDescriptionHelpFormatter, ) - # Required arguments - parser.add_argument("pid", type=int, help="Process ID to sample") + # Target selection + target_group = parser.add_mutually_exclusive_group(required=True) + target_group.add_argument( + "-p", "--pid", type=int, help="Process ID to sample" + ) + target_group.add_argument( + "-m", "--module", + nargs=argparse.REMAINDER, + help="Run and profile a module as python -m module [ARGS...]" + ) + target_group.add_argument( + "script", + nargs=argparse.REMAINDER, + help="Script to run and profile, with optional arguments" + ) # Sampling options sampling_group = parser.add_argument_group("Sampling configuration") @@ -712,19 +733,59 @@ def main(): sort_value = args.sort if args.sort is not None else 2 - sample( - args.pid, - sample_interval_usec=args.interval, - duration_sec=args.duration, - filename=args.outfile, - all_threads=args.all_threads, - limit=args.limit, - sort=sort_value, - show_summary=not args.no_summary, - output_format=args.format, - realtime_stats=args.realtime_stats, - ) + if not(args.pid or args.module or args.script): + parser.error( + "You must specify either a process ID (-p), a module (-m), or a script to run." + ) + if args.pid: + sample( + args.pid, + sample_interval_usec=args.interval, + duration_sec=args.duration, + filename=args.outfile, + all_threads=args.all_threads, + limit=args.limit, + sort=sort_value, + show_summary=not args.no_summary, + output_format=args.format, + realtime_stats=args.realtime_stats, + ) + elif args.module or args.script: + if args.module: + cmd = [sys.executable, "-m", *args.module] + else: + cmd = [sys.executable, *args.script] + + process = subprocess.Popen(cmd) + + try: + exit_code = process.wait(timeout=0.1) + sys.exit(exit_code) + except subprocess.TimeoutExpired: + pass + + try: + sample( + process.pid, + sort=sort_value, + sample_interval_usec=args.interval, + duration_sec=args.duration, + filename=args.outfile, + all_threads=args.all_threads, + limit=args.limit, + show_summary=not args.no_summary, + output_format=args.format, + realtime_stats=args.realtime_stats, + ) + finally: + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=2) + except subprocess.TimeoutExpired: + process.kill() + process.wait() if __name__ == "__main__": main() diff --git a/Lib/test/test_sample_profiler.py b/Lib/test/test_sample_profiler.py index 2c7fa1cba712c9..db11d520eaeef1 100644 --- a/Lib/test/test_sample_profiler.py +++ b/Lib/test/test_sample_profiler.py @@ -4,6 +4,7 @@ import io import marshal import os +import shutil import socket import subprocess import sys @@ -1588,6 +1589,68 @@ def test_sampling_all_threads(self): # Just verify that sampling completed without error # We're not testing output format here + def test_sample_target_script(self): + script_file = tempfile.NamedTemporaryFile(delete=False) + script_file.write(self.test_script.encode("utf-8")) + script_file.flush() + self.addCleanup(close_and_unlink, script_file) + + test_args = ["profile.sample", "-d", "1", script_file.name] + + with ( + mock.patch("sys.argv", test_args), + io.StringIO() as captured_output, + mock.patch("sys.stdout", captured_output), + ): + try: + profile.sample.main() + except PermissionError: + self.skipTest("Insufficient permissions for remote profiling") + + output = captured_output.getvalue() + + # Basic checks on output + self.assertIn("Captured", output) + self.assertIn("samples", output) + self.assertIn("Profile Stats", output) + + # Should see some of our test functions + self.assertIn("slow_fibonacci", output) + + + def test_sample_target_module(self): + tempdir = tempfile.TemporaryDirectory(delete=False) + self.addCleanup(lambda x: shutil.rmtree(x), tempdir.name) + + module_path = os.path.join(tempdir.name, "test_module.py") + + with open(module_path, "w") as f: + f.write(self.test_script) + + test_args = ["profile.sample", "-d", "1", "-m", "test_module"] + + with ( + mock.patch("sys.argv", test_args), + io.StringIO() as captured_output, + mock.patch("sys.stdout", captured_output), + # Change to temp directory so subprocess can find the module + contextlib.chdir(tempdir.name), + ): + try: + profile.sample.main() + except PermissionError: + self.skipTest("Insufficient permissions for remote profiling") + + output = captured_output.getvalue() + + # Basic checks on output + self.assertIn("Captured", output) + self.assertIn("samples", output) + self.assertIn("Profile Stats", output) + + # Should see some of our test functions + self.assertIn("slow_fibonacci", output) + @skip_if_not_supported @unittest.skipIf( @@ -1690,16 +1753,298 @@ def test_esrch_signal_handling(self): class TestSampleProfilerCLI(unittest.TestCase): + def test_cli_module_argument_parsing(self): + test_args = ["profile.sample", "-m", "mymodule"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profile.sample.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + ): + mock_process = mock.MagicMock() + mock_process.pid = 12345 + mock_process.wait.side_effect = [subprocess.TimeoutExpired(test_args, 0.1), None] + mock_process.poll.return_value = None + mock_popen.return_value = mock_process + + profile.sample.main() + + mock_popen.assert_called_once_with([sys.executable, "-m", "mymodule"]) + mock_sample.assert_called_once_with( + 12345, + sort=2, # default sort (sort_value from args.sort) + sample_interval_usec=100, + duration_sec=10, + filename=None, + all_threads=False, + limit=15, + show_summary=True, + output_format="pstats", + realtime_stats=False, + ) + + def test_cli_module_with_arguments(self): + test_args = ["profile.sample", "-m", "mymodule", "arg1", "arg2", "--flag"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profile.sample.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + ): + mock_process = mock.MagicMock() + mock_process.pid = 12345 + mock_process.wait.side_effect = [subprocess.TimeoutExpired(test_args, 0.1), None] + mock_process.poll.return_value = None + mock_popen.return_value = mock_process + + profile.sample.main() + + mock_popen.assert_called_once_with([ + sys.executable, "-m", "mymodule", "arg1", "arg2", "--flag" + ]) + mock_sample.assert_called_once_with( + 12345, + sort=2, + sample_interval_usec=100, + duration_sec=10, + filename=None, + all_threads=False, + limit=15, + show_summary=True, + output_format="pstats", + realtime_stats=False, + ) + + def test_cli_script_argument_parsing(self): + test_args = ["profile.sample", "myscript.py"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profile.sample.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + ): + mock_process = mock.MagicMock() + mock_process.pid = 12345 + mock_process.wait.side_effect = [subprocess.TimeoutExpired(test_args, 0.1), None] + mock_process.poll.return_value = None + mock_popen.return_value = mock_process + + profile.sample.main() + + mock_popen.assert_called_once_with([sys.executable, "myscript.py"]) + mock_sample.assert_called_once_with( + 12345, + sort=2, + sample_interval_usec=100, + duration_sec=10, + filename=None, + all_threads=False, + limit=15, + show_summary=True, + output_format="pstats", + realtime_stats=False, + ) + + def test_cli_script_with_arguments(self): + test_args = ["profile.sample", "myscript.py", "arg1", "arg2", "--flag"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profile.sample.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + ): + mock_process = mock.MagicMock() + mock_process.pid = 12345 + mock_process.wait.side_effect = [subprocess.TimeoutExpired(test_args, 0.1), None] + mock_process.poll.return_value = None + mock_popen.return_value = mock_process + + profile.sample.main() + + mock_popen.assert_called_once_with([ + sys.executable, "myscript.py", "arg1", "arg2", "--flag" + ]) + + def test_cli_mutually_exclusive_pid_module(self): + test_args = ["profile.sample", "-p", "12345", "-m", "mymodule"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + profile.sample.main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("not allowed with argument", error_msg) + + def test_cli_mutually_exclusive_pid_script(self): + test_args = ["profile.sample", "-p", "12345", "myscript.py"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + profile.sample.main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("not allowed with argument", error_msg) + + def test_cli_no_target_specified(self): + test_args = ["profile.sample", "-d", "5"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + profile.sample.main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("one of the arguments", error_msg) + + def test_cli_module_with_profiler_options(self): + test_args = [ + "profile.sample", "-i", "1000", "-d", "30", "-a", + "--sort-tottime", "-l", "20", "-m", "mymodule", + ] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profile.sample.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + ): + mock_process = mock.MagicMock() + mock_process.pid = 12345 + mock_process.wait.side_effect = [subprocess.TimeoutExpired(test_args, 0.1), None] + mock_process.poll.return_value = None + mock_popen.return_value = mock_process + + profile.sample.main() + + mock_sample.assert_called_once_with( + 12345, + sort=1, # sort-tottime + sample_interval_usec=1000, + duration_sec=30, + filename=None, + all_threads=True, + limit=20, + show_summary=True, + output_format="pstats", + realtime_stats=False, + ) + + def test_cli_script_with_profiler_options(self): + """Test script with various profiler options.""" + test_args = [ + "profile.sample", "-i", "2000", "-d", "60", + "--collapsed", "-o", "output.txt", + "myscript.py", "scriptarg", + ] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profile.sample.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + ): + mock_process = mock.MagicMock() + mock_process.pid = 12345 + mock_process.wait.side_effect = [subprocess.TimeoutExpired(test_args, 0.1), None] + mock_process.poll.return_value = None + mock_popen.return_value = mock_process + + profile.sample.main() + + mock_popen.assert_called_once_with([ + sys.executable, "myscript.py", "scriptarg" + ]) + # Verify profiler options were passed correctly + mock_sample.assert_called_once_with( + 12345, + sort=2, # default sort + sample_interval_usec=2000, + duration_sec=60, + filename="output.txt", + all_threads=False, + limit=15, + show_summary=True, + output_format="collapsed", + realtime_stats=False, + ) + + def test_cli_empty_module_name(self): + test_args = ["profile.sample", "-m"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("sys.stderr", io.StringIO()) as mock_stderr, + self.assertRaises(SystemExit) as cm, + ): + profile.sample.main() + + self.assertEqual(cm.exception.code, 2) # argparse error + error_msg = mock_stderr.getvalue() + self.assertIn("must specify", error_msg) + + def test_cli_long_module_option(self): + test_args = ["profile.sample", "--module", "mymodule", "arg1"] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profile.sample.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + ): + mock_process = mock.MagicMock() + mock_process.pid = 12345 + mock_process.wait.side_effect = [subprocess.TimeoutExpired(test_args, 0.1), None] + mock_process.poll.return_value = None + mock_popen.return_value = mock_process + + profile.sample.main() + + mock_popen.assert_called_once_with([ + sys.executable, "-m", "mymodule", "arg1", + ]) + + def test_cli_complex_script_arguments(self): + test_args = [ + "profile.sample", "script.py", + "--input", "file.txt", "-v", "--output=/tmp/out", "positional" + ] + + with ( + mock.patch("sys.argv", test_args), + mock.patch("profile.sample.sample") as mock_sample, + mock.patch("subprocess.Popen") as mock_popen, + ): + mock_process = mock.MagicMock() + mock_process.pid = 12345 + mock_process.wait.side_effect = [subprocess.TimeoutExpired(test_args, 0.1), None] + mock_process.poll.return_value = None + mock_popen.return_value = mock_process + + profile.sample.main() + + mock_popen.assert_called_once_with([ + sys.executable, "script.py", + "--input", "file.txt", "-v", "--output=/tmp/out", "positional", + ]) + def test_cli_collapsed_format_validation(self): """Test that CLI properly validates incompatible options with collapsed format.""" test_cases = [ # Test sort options are invalid with collapsed ( - ["profile.sample", "--collapsed", "--sort-nsamples", "12345"], + ["profile.sample", "--collapsed", "--sort-nsamples", "-p", "12345"], "sort", ), ( - ["profile.sample", "--collapsed", "--sort-tottime", "12345"], + ["profile.sample", "--collapsed", "--sort-tottime", "-p", "12345"], "sort", ), ( @@ -1707,6 +2052,7 @@ def test_cli_collapsed_format_validation(self): "profile.sample", "--collapsed", "--sort-cumtime", + "-p", "12345", ], "sort", @@ -1716,6 +2062,7 @@ def test_cli_collapsed_format_validation(self): "profile.sample", "--collapsed", "--sort-sample-pct", + "-p", "12345", ], "sort", @@ -1725,23 +2072,24 @@ def test_cli_collapsed_format_validation(self): "profile.sample", "--collapsed", "--sort-cumul-pct", + "-p", "12345", ], "sort", ), ( - ["profile.sample", "--collapsed", "--sort-name", "12345"], + ["profile.sample", "--collapsed", "--sort-name", "-p", "12345"], "sort", ), # Test limit option is invalid with collapsed - (["profile.sample", "--collapsed", "-l", "20", "12345"], "limit"), + (["profile.sample", "--collapsed", "-l", "20", "-p", "12345"], "limit"), ( - ["profile.sample", "--collapsed", "--limit", "20", "12345"], + ["profile.sample", "--collapsed", "--limit", "20", "-p", "12345"], "limit", ), # Test no-summary option is invalid with collapsed ( - ["profile.sample", "--collapsed", "--no-summary", "12345"], + ["profile.sample", "--collapsed", "--no-summary", "-p", "12345"], "summary", ), ] @@ -1761,7 +2109,7 @@ def test_cli_collapsed_format_validation(self): def test_cli_default_collapsed_filename(self): """Test that collapsed format gets a default filename when not specified.""" - test_args = ["profile.sample", "--collapsed", "12345"] + test_args = ["profile.sample", "--collapsed", "-p", "12345"] with ( mock.patch("sys.argv", test_args), @@ -1779,12 +2127,12 @@ def test_cli_custom_output_filenames(self): """Test custom output filenames for both formats.""" test_cases = [ ( - ["profile.sample", "--pstats", "-o", "custom.pstats", "12345"], + ["profile.sample", "--pstats", "-o", "custom.pstats", "-p", "12345"], "custom.pstats", "pstats", ), ( - ["profile.sample", "--collapsed", "-o", "custom.txt", "12345"], + ["profile.sample", "--collapsed", "-o", "custom.txt", "-p", "12345"], "custom.txt", "collapsed", ), @@ -1816,7 +2164,7 @@ def test_cli_mutually_exclusive_format_options(self): with ( mock.patch( "sys.argv", - ["profile.sample", "--pstats", "--collapsed", "12345"], + ["profile.sample", "--pstats", "--collapsed", "-p", "12345"], ), mock.patch("sys.stderr", io.StringIO()), ): @@ -1824,7 +2172,7 @@ def test_cli_mutually_exclusive_format_options(self): profile.sample.main() def test_argument_parsing_basic(self): - test_args = ["profile.sample", "12345"] + test_args = ["profile.sample", "-p", "12345"] with ( mock.patch("sys.argv", test_args), @@ -1856,7 +2204,7 @@ def test_sort_options(self): ] for option, expected_sort_value in sort_options: - test_args = ["profile.sample", option, "12345"] + test_args = ["profile.sample", option, "-p", "12345"] with ( mock.patch("sys.argv", test_args),
Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.
Alternative Proxies: