diff --git a/psrecord/main.py b/psrecord/main.py index eb8385b..a4f07e2 100644 --- a/psrecord/main.py +++ b/psrecord/main.py @@ -30,6 +30,10 @@ children = [] +# Characters that mean a command has to be run through a shell rather than +# being split into arguments and executed directly +SHELL_CHARS = "|&;<>()$`\\\"'*?[]{}~#\n" + def get_percent(process): return process.cpu_percent() @@ -102,17 +106,32 @@ def main(): args = parser.parse_args() + include_children = args.include_children + # Attach to process try: pid = int(args.process_id_or_command) print(f"Attaching to process {pid}") sprocess = None except Exception: + import shlex import subprocess command = args.process_id_or_command print(f"Starting up command '{command}' and attaching to process") - sprocess = subprocess.Popen(command, shell=True) + if sys.platform == "win32": + sprocess = subprocess.Popen(command) + elif any(char in command for char in SHELL_CHARS): + # Running the command through a shell means the monitored process + # is the shell itself, since modern shells no longer exec single + # commands in-place, so child processes have to be included for + # the statistics to be meaningful. + sprocess = subprocess.Popen(command, shell=True) + if not include_children: + print("Command requires a shell, including child processes in statistics") + include_children = True + else: + sprocess = subprocess.Popen(shlex.split(command)) pid = sprocess.pid monitor( @@ -121,7 +140,7 @@ def main(): plot=args.plot, duration=args.duration, interval=args.interval, - include_children=args.include_children, + include_children=include_children, include_io=args.include_io, log_format=args.log_format, ) diff --git a/psrecord/tests/test_main.py b/psrecord/tests/test_main.py index 9e0037d..8bf313b 100644 --- a/psrecord/tests/test_main.py +++ b/psrecord/tests/test_main.py @@ -72,7 +72,11 @@ def test_plot(self, tmpdir): assert os.path.exists(filename) def test_main(self): - sys.argv = ["psrecord", "--duration=3", "'sleep 10'"] + sys.argv = [ + "psrecord", + "--duration=3", + f'"{sys.executable}" -c "import time; time.sleep(10)"', + ] main() def test_main_by_id(self): @@ -82,3 +86,60 @@ def test_main_by_id(self): @pytest.mark.skipif(sys.platform == "darwin", reason="Functionality not supported on MacOS") def test_io(self, tmpdir): monitor(os.getpid(), duration=3, include_io=True) + + +MEMORY_CODE = """ +x = [0] * 20_000_000 +import time +time.sleep(10) +""" + + +def max_logged_memory(filename): + with open(filename) as f: + rows = [line.split() for line in f.readlines()[1:]] + return max(float(row[2]) for row in rows) + + +@pytest.mark.skipif(sys.platform == "win32", reason="Test relies on POSIX commands") +class TestLaunchedCommand: + # When psrecord launches the command itself, the logged statistics + # should be those of the command, not of an intermediate shell + + def setup_method(self, method): + self.original_argv = sys.argv + + def teardown_method(self, method): + sys.argv = self.original_argv + + def test_simple_command(self, tmpdir): + script = tmpdir.join("memory.py").strpath + with open(script, "w") as f: + f.write(MEMORY_CODE) + logfile = tmpdir.join("log.txt").strpath + sys.argv = [ + "psrecord", + f"{sys.executable} {script}", + "--log", + logfile, + "--duration=3", + "--interval=0.2", + ] + main() + assert max_logged_memory(logfile) > 100 + + def test_command_with_shell_syntax(self, tmpdir): + script = tmpdir.join("memory.py").strpath + with open(script, "w") as f: + f.write(MEMORY_CODE) + logfile = tmpdir.join("log.txt").strpath + sys.argv = [ + "psrecord", + f"sleep 0 && {sys.executable} {script}", + "--log", + logfile, + "--duration=3", + "--interval=0.2", + ] + main() + assert max_logged_memory(logfile) > 100