Repository: trailofbits/pip-audit (python)

Issue:
Make `pip-audit`'s spinner still more responsive
This is not triaged/prioritized since it's just a braindump.

#56 and #89 made `pip-audit` more responsive, but we can still do better.

In particular, the spinner still pauses when `pip-audit` is performing any kind of subprocess (i.e., calling out to `pip` internally).

We can probably improve this by adding some kind of synchronous monitoring of each subprocess -- we can either update the spinner naively (i.e., on a cadence, once every `X` milliseconds while waiting for the subprocess to terminate), or by monitoring the process's output and updating the spinner e.g. each time we see a new line on `stdout` or `stderr`.

Upsides:

* The `pip-audit` CLI becomes even more responsive, particularly in subprocess-heavy modes like `-r` and `--fix`

Downsides:

* We probably have to use one of the messier `subprocess` APIs, like `subprocess.communicate`

Pull request description:
Make the CLI more responsive
Closes #227.

Demo:

[![asciicast](https://asciinema.org/a/465724.svg)](https://asciinema.org/a/465724)

Patch:
diff --git a/pip_audit/_subprocess.py b/pip_audit/_subprocess.py
new file mode 100644
index 0000000..8692a1c
--- /dev/null
+++ b/pip_audit/_subprocess.py
@@ -0,0 +1,53 @@
+"""
+A thin `subprocess` wrapper for making long-running subprocesses more
+responsive from the `pip-audit` CLI.
+"""
+
+import os.path
+import subprocess
+from subprocess import Popen
+from typing import Sequence
+
+from ._state import AuditState
+
+
+class CalledProcessError(Exception):
+    """
+    Raised if the underlying subprocess created by `run` exits with a nonzero code.
+    """
+
+
+def run(args: Sequence[str], *, state: AuditState = AuditState()) -> str:
+    """
+    Execute the given arguments.
+
+    Uses `state` to provide feedback on the subprocess's status.
+
+    Raises a `CalledProcessError` if the subprocess fails. Otherwise, returns
+    the process's `stdout` stream as a string.
+    """
+
+    # Run the process with unbuffered I/O, to make the poll-and-read loop below
+    # more responsive.
+    process = Popen(args, bufsize=0, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
+
+    # NOTE(ww): We frequently run commands inside of ephemeral virtual environments,
+    # which have long absolute paths on some platforms. These make for confusing
+    # state updates, so we trim the first argument down to its basename.
+    pretty_args = " ".join([os.path.basename(args[0]), *args[1:]])
+
+    terminated = False
+    stdout = bytes()
+
+    # NOTE: We use `poll()` to control this loop instead of the `read()` call
+    # to prevent deadlocks. Similarly, `read(size)` will return an empty bytes
+    # once `stdout` hits EOF, so we don't have to worry about that blocking.
+    while not terminated:
+        terminated = process.poll() is not None
+        stdout += process.stdout.read(4096)  # type: ignore
+        state.update_state(f"Running {pretty_args}")
+
+    if process.returncode != 0:
+        raise CalledProcessError(f"{pretty_args} exited with {process.returncode}")
+
+    return stdout.decode("utf-8")
diff --git a/pip_audit/_virtual_env.py b/pip_audit/_virtual_env.py
index 52d9a09..713810e 100644
--- a/pip_audit/_virtual_env.py
+++ b/pip_audit/_virtual_env.py
@@ -4,13 +4,13 @@ Create virtual environments with a custom set of packages and inspect their depe
 
 import json
 import logging
-import subprocess
 import venv
 from typing import Iterator, List, Optional, Tuple
 
 from packaging.version import Version
 
 from ._state import AuditState
+from ._subprocess import CalledProcessError, run
 
 logger = logging.getLogger(__name__)
 
@@ -85,10 +85,8 @@ class VirtualEnv(venv.EnvBuilder):
             "wheel",
         ]
         try:
-            subprocess.run(
-                pip_upgrade_cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
-            )
-        except subprocess.CalledProcessError as cpe:
+            run(pip_upgrade_cmd, state=self._state)
+        except CalledProcessError as cpe:
             raise VirtualEnvError(f"Failed to upgrade `pip`: {pip_upgrade_cmd}") from cpe
 
         self._state.update_state("Installing package in isolated environment")
@@ -102,13 +100,8 @@ class VirtualEnv(venv.EnvBuilder):
             *self._install_args,
         ]
         try:
-            subprocess.run(
-                package_install_cmd,
-                check=True,
-                stdout=subprocess.DEVNULL,
-                stderr=subprocess.DEVNULL,
-            )
-        except subprocess.CalledProcessError as cpe:
+            run(package_install_cmd, state=self._state)
+        except CalledProcessError as cpe:
             raise VirtualEnvError(f"Failed to install packages: {package_install_cmd}") from cpe
 
         self._state.update_state("Processing package list from isolated environment")
@@ -117,13 +110,10 @@ class VirtualEnv(venv.EnvBuilder):
         # environment contains
         list_cmd = [context.env_exe, "-m", "pip", "list", "-l", "--format", "json"]
         try:
-            process = subprocess.run(
-                list_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL
-            )
-        except subprocess.CalledProcessError as cpe:
+            stdout = run(list_cmd, state=self._state)
+        except CalledProcessError as cpe:
             raise VirtualEnvError(f"Failed to run `pip list`: {list_cmd}") from cpe
-        list_output = process.stdout.decode("utf-8")
-        package_list = json.loads(list_output)
+        package_list = json.loads(stdout)
 
         # Convert into a series of name, version pairs
         self._packages = []


Test patch:
diff --git a/test/test_subprocess.py b/test/test_subprocess.py
new file mode 100644
index 0000000..21d9ccb
--- /dev/null
+++ b/test/test_subprocess.py
@@ -0,0 +1,8 @@
+import pytest
+
+from pip_audit._subprocess import CalledProcessError, run
+
+
+def test_run_raises():
+    with pytest.raises(CalledProcessError):
+        run(["false"])
diff --git a/test/test_virtual_env.py b/test/test_virtual_env.py
index 76e32f7..12c6ccf 100644
--- a/test/test_virtual_env.py
+++ b/test/test_virtual_env.py
@@ -1,9 +1,10 @@
-import subprocess
 from tempfile import TemporaryDirectory
 
 import pytest
 from packaging.version import Version
 
+import pip_audit._virtual_env as _virtual_env
+from pip_audit import _subprocess
 from pip_audit._virtual_env import VirtualEnv, VirtualEnvError
 
 
@@ -23,15 +24,15 @@ def test_virtual_env_incorrect_usage():
 
 
 def test_virtual_env_failed_package_installation(monkeypatch):
-    original_run = subprocess.run
+    original_run = _subprocess.run
 
     def run_mock(args, **kwargs):
         if "flask==2.0.1" in args:
-            raise subprocess.CalledProcessError(-1, str())
+            raise _subprocess.CalledProcessError("barf")
         # If it's not the package installation command, then call the original run
         return original_run(args, **kwargs)
 
-    monkeypatch.setattr(subprocess, "run", run_mock)
+    monkeypatch.setattr(_virtual_env, "run", run_mock)
 
     with TemporaryDirectory() as ve_dir:
         ve = VirtualEnv(["flask==2.0.1"])
@@ -40,17 +41,17 @@ def test_virtual_env_failed_package_installation(monkeypatch):
 
 
 def test_virtual_env_failed_pip_upgrade(monkeypatch):
-    original_run = subprocess.run
+    original_run = _subprocess.run
 
     def run_mock(args, **kwargs):
         # We have to be a bit more specific than usual here because the `EnvBuilder` invokes
         # `ensurepip` with similar looking arguments and we DON'T want to mock that call.
         if set(["install", "--upgrade", "pip"]).issubset(set(args)):
-            raise subprocess.CalledProcessError(-1, str())
+            raise _subprocess.CalledProcessError("barf")
         # If it's not a call to upgrade pip, then call the original run
         return original_run(args, **kwargs)
 
-    monkeypatch.setattr(subprocess, "run", run_mock)
+    monkeypatch.setattr(_virtual_env, "run", run_mock)
 
     with TemporaryDirectory() as ve_dir:
         ve = VirtualEnv(["flask==2.0.1"])
@@ -59,15 +60,15 @@ def test_virtual_env_failed_pip_upgrade(monkeypatch):
 
 
 def test_virtual_env_failed_pip_list(monkeypatch):
-    original_run = subprocess.run
+    original_run = _subprocess.run
 
     def run_mock(args, **kwargs):
-        if "json" in args:
-            raise subprocess.CalledProcessError(-1, str())
+        if {"list", "--format", "json"}.issubset(set(args)):
+            raise _subprocess.CalledProcessError("barf")
         # If it's not a call to `pip list`, then call the original run
         return original_run(args, **kwargs)
 
-    monkeypatch.setattr(subprocess, "run", run_mock)
+    monkeypatch.setattr(_virtual_env, "run", run_mock)
 
     with TemporaryDirectory() as ve_dir:
         ve = VirtualEnv(["flask==2.0.1"])