diff --git a/CHANGES.md b/CHANGES.md index 51f74d6..ee0d50e 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -15,6 +15,8 @@ for contemporary versions of Java - Add support for Python 3.7 - Add `testbook` to `notebook` subsystem +- Add `pueblo.util.proc.process` utility. It is a wrapper around + `subprocess.Popen` to also terminate child processes after exiting. ## 2023-11-06 v0.0.3 - ngr: Fix `contextlib.chdir` only available on Python 3.11 and newer diff --git a/pueblo/util/proc.py b/pueblo/util/proc.py new file mode 100644 index 0000000..fa2ea6d --- /dev/null +++ b/pueblo/util/proc.py @@ -0,0 +1,26 @@ +import subprocess +import typing as t +from contextlib import contextmanager + +import psutil + + +@contextmanager +def process(*args, **kwargs) -> t.Generator[subprocess.Popen, None, None]: + """ + Wrapper around `subprocess.Popen` to also terminate child processes after exiting. + + Implementation by Pedro Cattori. Thanks! + -- https://gist.github.com/jizhilong/6687481#gistcomment-3057122 + """ + proc = subprocess.Popen(*args, **kwargs) # noqa: S603 + try: + yield proc + finally: + try: + children = psutil.Process(proc.pid).children(recursive=True) + except psutil.NoSuchProcess: + return + for child in children: + child.kill() + proc.kill() diff --git a/pyproject.toml b/pyproject.toml index c75d6c9..eecf254 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,7 +70,7 @@ dependencies = [ [project.optional-dependencies] all = [ - "pueblo[cli,dataframe,fileio,nlp,notebook,testing,web]", + "pueblo[cli,dataframe,fileio,nlp,notebook,proc,testing,web]", ] cli = [ "click<9", @@ -109,6 +109,9 @@ notebook = [ "pytest-notebook<0.11", "testbook<0.5", ] +proc = [ + "psutil<6", +] release = [ "build<2", "twine<5", diff --git a/tests/test_proc.py b/tests/test_proc.py new file mode 100644 index 0000000..0de480f --- /dev/null +++ b/tests/test_proc.py @@ -0,0 +1,32 @@ +import pytest + +from pueblo.util.proc import process + + +def test_process_success(tmp_path): + outfile = tmp_path / "myfile.out.log" + with process(["echo", "Hello, world."], stdout=open(outfile, "w")) as proc: + assert isinstance(proc.pid, int) + + with open(outfile, "r") as fp: + assert fp.read() == "Hello, world.\n" + + +def test_process_noop(): + process(["mycommand", "myarg", "--myoption", "myoptionvalue"]) + + +def test_process_failure_command_not_found(): + with pytest.raises(FileNotFoundError) as ex: + with process(["mycommand", "myarg", "--myoption", "myoptionvalue"]): + pass + assert ex.match("No such file or directory") + + +def test_process_failure_in_contextmanager(): + with pytest.raises(ZeroDivisionError): + with process(["echo", "Hello, world."]) as proc: + print(proc.pid) # noqa: T201 + # Even though this throws an exception, the `process` contextmanager + # will *still* clean up the process correctly. + 0 / 0 # noqa: B018