""" Helper functions for managing processes. """ from __future__ import print_function import sys import os import subprocess import signal import psutil import atexit def kill_process(proc): """ Kill the process `proc` created with `subprocess`. """ p1_group = psutil.Process(proc.pid) child_pids = p1_group.get_children(recursive=True) for child_pid in child_pids: os.kill(child_pid.pid, signal.SIGKILL) def run_multi_processes(cmd_list, out_log=None, err_log=None): """ Run each shell command in `cmd_list` in a separate process, piping stdout to `out_log` (a path) and stderr to `err_log` (also a path). Terminates the processes on CTRL-C and ensures the processes are killed if an error occurs. """ kwargs = {'shell': True, 'cwd': None} pids = [] if out_log: out_log_file = open(out_log, 'w') kwargs['stdout'] = out_log_file if err_log: err_log_file = open(err_log, 'w') kwargs['stderr'] = err_log_file try: for cmd in cmd_list: pids.extend([subprocess.Popen(cmd, **kwargs)]) def _signal_handler(*args): """ What to do when process is ended """ print("\nEnding...") signal.signal(signal.SIGINT, _signal_handler) print("Enter CTL-C to end") signal.pause() print("Processes ending") except Exception as err: print("Error running process {}".format(err), file=sys.stderr) finally: for pid in pids: kill_process(pid) def run_process(cmd, out_log=None, err_log=None): """ Run the shell command `cmd` in a separate process, piping stdout to `out_log` (a path) and stderr to `err_log` (also a path). Terminates the process on CTRL-C or if an error occurs. """ return run_multi_processes([cmd], out_log=out_log, err_log=err_log) def run_background_process(cmd, out_log=None, err_log=None, cwd=None): """ Runs a command as a background process. Sends SIGINT at exit. """ kwargs = {'shell': True, 'cwd': cwd} if out_log: out_log_file = open(out_log, 'w') kwargs['stdout'] = out_log_file if err_log: err_log_file = open(err_log, 'w') kwargs['stderr'] = err_log_file proc = subprocess.Popen(cmd, **kwargs) def exit_handler(): """ Send SIGINT to the process's children. This is important for running commands under coverage, as coverage will not produce the correct artifacts if the child process isn't killed properly. """ p1_group = psutil.Process(proc.pid) child_pids = p1_group.get_children(recursive=True) for child_pid in child_pids: os.kill(child_pid.pid, signal.SIGINT) atexit.register(exit_handler)