Repository: trailofbits/pip-audit (python)

Issue:
Make the CLI *even more* responsive
Follow-up on #56, per https://github.com/trailofbits/pip-audit/pull/87#issuecomment-954320210:

> Regarding the feedback not being great at the beginning, in the future it might be worth reporting some extra state in the PyPIProvider. Something like:

> `Gathering metadata for sdist {name} ({version})`

Pull request:
pip_audit: Update state feedback when fetching metadata
Closes #89

- [x] Update state for virtual env creation.
- [x] Update state for sdist installation

Patch:
diff --git a/pip_audit/dependency_source/resolvelib/pypi_provider.py b/pip_audit/dependency_source/resolvelib/pypi_provider.py
index 65c999d..34b5a37 100644
--- a/pip_audit/dependency_source/resolvelib/pypi_provider.py
+++ b/pip_audit/dependency_source/resolvelib/pypi_provider.py
@@ -32,12 +32,21 @@ PYTHON_VERSION = Version(python_version())
 
 
 class Candidate:
-    def __init__(self, name, version, url=None, extras=None, is_wheel=True):
+    def __init__(
+        self,
+        name,
+        version,
+        url=None,
+        extras=None,
+        is_wheel=True,
+        state: Optional[AuditState] = None,
+    ):
         self.name = canonicalize_name(name)
         self.version = version
         self.url = url
         self.extras = extras
         self.is_wheel = is_wheel
+        self.state = state
 
         self._metadata = None
         self._dependencies = None
@@ -50,10 +59,13 @@ class Candidate:
     @property
     def metadata(self):
         if self._metadata is None:
+            if self.state is not None:  # pragma: no cover
+                self.state.update_state(f"Fetching metadata for {self.name} ({self.version})")
+
             if self.is_wheel:
-                self._metadata = get_metadata_for_wheel(self.url)
+                self._metadata = self._get_metadata_for_wheel()
             else:
-                self._metadata = get_metadata_for_sdist(self.url)
+                self._metadata = self._get_metadata_for_sdist()
         return self._metadata
 
     def _get_dependencies(self):
@@ -75,8 +87,68 @@ class Candidate:
             self._dependencies = list(self._get_dependencies())
         return self._dependencies
 
+    def _get_metadata_for_wheel(self):
+        data = requests.get(self.url).content
 
-def get_project_from_pypi(project, extras):
+        if self.state is not None:
+            self.state.update_state(
+                f"Extracting wheel for {self.name} ({self.version})"
+            )  # pragma: no cover
+
+        with ZipFile(BytesIO(data)) as z:
+            for n in z.namelist():
+                if n.endswith(".dist-info/METADATA"):
+                    p = BytesParser()
+                    return p.parse(z.open(n), headersonly=True)
+
+        # If we didn't find the metadata, return an empty dict
+        return EmailMessage()  # pragma: no cover
+
+    def _get_metadata_for_sdist(self):
+        response: requests.Response = requests.get(self.url)
+        response.raise_for_status()
+        data = response.content
+        metadata = EmailMessage()
+
+        with TemporaryDirectory() as pkg_dir:
+            if self.state is not None:
+                self.state.update_state(
+                    f"Extracting source distribution for {self.name} ({self.version})"
+                )  # pragma: no cover
+
+            # Extract archive onto the disk
+            with TarFile.open(fileobj=BytesIO(data), mode="r:gz") as t:
+                # The directory is the first member in a tarball
+                names = t.getnames()
+                pkg_name = names[0]
+                t.extractall(pkg_dir)
+
+            if self.state is not None:
+                self.state.update_state(
+                    f"Installing source distribution in isolated environment for {self.name} "
+                    f"({self.version})"
+                )  # pragma: no cover
+
+            # Put together a full path of where the source distribution is
+            pkg_path = os.path.join(pkg_dir, pkg_name)
+
+            with TemporaryDirectory() as ve_dir:
+                ve = VirtualEnv(["-e", pkg_path], self.state)
+                ve.create(ve_dir)
+
+                if self.state is not None:
+                    self.state.update_state(
+                        f"Querying installed packages for {self.name} ({self.version})"
+                    )  # pragma: no cover
+
+                installed_packages = ve.installed_packages
+                for name, version in installed_packages:
+                    metadata["Requires-Dist"] = f"{name}=={str(version)}"
+
+        return metadata
+
+
+def get_project_from_pypi(project, extras, state: Optional[AuditState]):
     """Return candidates created from the project name and extras."""
     url = "https://pypi.org/simple/{}".format(project)
     response: requests.Response = requests.get(url)
@@ -111,45 +183,7 @@ def get_project_from_pypi(project, extras):
 
         # TODO: Handle compatibility tags?
 
-        yield Candidate(name, version, url=url, extras=extras, is_wheel=is_wheel)
-
-
-def get_metadata_for_wheel(url):
-    data = requests.get(url).content
-    with ZipFile(BytesIO(data)) as z:
-        for n in z.namelist():
-            if n.endswith(".dist-info/METADATA"):
-                p = BytesParser()
-                return p.parse(z.open(n), headersonly=True)
-
-    # If we didn't find the metadata, return an empty dict
-    return EmailMessage()  # pragma: no cover
-
-
-def get_metadata_for_sdist(url):
-    response: requests.Response = requests.get(url)
-    response.raise_for_status()
-    data = response.content
-    metadata = EmailMessage()
-
-    with TemporaryDirectory() as pkg_dir:
-        # Extract archive onto the disk
-        with TarFile.open(fileobj=BytesIO(data), mode="r:gz") as t:
-            # The directory is the first member in a tarball
-            names = t.getnames()
-            pkg_name = names[0]
-            t.extractall(pkg_dir)
-
-        # Put together a full path of where the source distribution is
-        pkg_path = os.path.join(pkg_dir, pkg_name)
-
-        with TemporaryDirectory() as ve_dir:
-            ve = VirtualEnv(["-e", pkg_path])
-            ve.create(ve_dir)
-            for name, version in ve.installed_packages:
-                metadata["Requires-Dist"] = f"{name}=={str(version)}"
-
-    return metadata
+        yield Candidate(name, version, url=url, extras=extras, is_wheel=is_wheel, state=state)
 
 
 class PyPIProvider(AbstractProvider):
@@ -180,7 +214,7 @@ class PyPIProvider(AbstractProvider):
         # treat candidates as immutable once created.
         candidates = (
             candidate
-            for candidate in get_project_from_pypi(identifier, extras)
+            for candidate in get_project_from_pypi(identifier, extras, self.state)
             if candidate.version not in bad_versions
             and all(candidate.version in r.specifier for r in requirements)
         )
diff --git a/pip_audit/virtual_env.py b/pip_audit/virtual_env.py
index f104d10..9ed52a9 100644
--- a/pip_audit/virtual_env.py
+++ b/pip_audit/virtual_env.py
@@ -9,15 +9,23 @@ from typing import Iterator, List, Optional, Tuple
 
 from packaging.version import Version
 
+from .state import AuditState
+
 
 class VirtualEnv(venv.EnvBuilder):
-    def __init__(self, install_args: List[str]):
+    def __init__(self, install_args: List[str], state: Optional[AuditState] = None):
         super().__init__(with_pip=True)
         self._install_args = install_args
         self._packages: Optional[List[Tuple[str, Version]]] = None
+        self._state = state
 
     # Override this hook with custom behaviour
     def post_setup(self, context):
+        if self._state is not None:
+            self._state.update_state(
+                "Updating pip installation in isolated environment"
+            )  # pragma: no cover
+
         # Firstly, upgrade our `pip` versions since `ensurepip` can leave us with an old version
         pip_upgrade_cmd = [context.env_exe, "-m", "pip", "install", "--upgrade", "pip"]
         try:
@@ -27,6 +35,11 @@ class VirtualEnv(venv.EnvBuilder):
         except subprocess.CalledProcessError as cpe:
             raise VirtualEnvError(f"Failed to upgrade `pip`: {pip_upgrade_cmd}") from cpe
 
+        if self._state is not None:
+            self._state.update_state(
+                "Installing package in isolated environment"
+            )  # pragma: no cover
+
         # Install our packages
         package_install_cmd = [context.env_exe, "-m", "pip", "install", *self._install_args]
         try:
@@ -39,6 +52,11 @@ class VirtualEnv(venv.EnvBuilder):
         except subprocess.CalledProcessError as cpe:
             raise VirtualEnvError(f"Failed to install packages: {package_install_cmd}") from cpe
 
+        if self._state is not None:
+            self._state.update_state(
+                "Processing package list from isolated environment"
+            )  # pragma: no cover
+
         # Now parse the `pip list` output to figure out what packages our
         # environment contains
         list_cmd = [context.env_exe, "-m", "pip", "list", "-l", "--format", "json"]


Test patch:
diff --git a/test/dependency_source/test_resolvelib.py b/test/dependency_source/test_resolvelib.py
index 4511ba5..fff04d6 100644
--- a/test/dependency_source/test_resolvelib.py
+++ b/test/dependency_source/test_resolvelib.py
@@ -126,7 +126,9 @@ def test_resolvelib_wheel_patched(monkeypatch):
     )
 
     monkeypatch.setattr(requests, "get", lambda _: get_package_mock(data))
-    monkeypatch.setattr(pypi_provider, "get_metadata_for_wheel", lambda _: get_metadata_mock())
+    monkeypatch.setattr(
+        pypi_provider.Candidate, "_get_metadata_for_wheel", lambda _: get_metadata_mock()
+    )
 
     resolver = resolvelib.ResolveLibResolver()
     req = Requirement("flask==2.0.1")
@@ -144,7 +146,9 @@ def test_resolvelib_sdist_patched(monkeypatch, suffix):
     data = f'<a href="https://example.com/Flask-2.0.1.{suffix}">Flask-2.0.1.{suffix}</a><br/>'
 
     monkeypatch.setattr(requests, "get", lambda _: get_package_mock(data))
-    monkeypatch.setattr(pypi_provider, "get_metadata_for_sdist", lambda _: get_metadata_mock())
+    monkeypatch.setattr(
+        pypi_provider.Candidate, "_get_metadata_for_sdist", lambda _: get_metadata_mock()
+    )
 
     resolver = resolvelib.ResolveLibResolver()
     req = Requirement("flask==2.0.1")
@@ -183,7 +187,9 @@ def test_resolvelib_wheel_canonical_name_mismatch(monkeypatch):
     )
 
     monkeypatch.setattr(requests, "get", lambda _: get_package_mock(data))
-    monkeypatch.setattr(pypi_provider, "get_metadata_for_wheel", lambda _: get_metadata_mock())
+    monkeypatch.setattr(
+        pypi_provider.Candidate, "_get_metadata_for_wheel", lambda _: get_metadata_mock()
+    )
 
     resolver = resolvelib.ResolveLibResolver()
     req = Requirement("flask==2.0.1")
@@ -202,7 +208,9 @@ def test_resolvelib_wheel_invalid_version(monkeypatch):
     )
 
     monkeypatch.setattr(requests, "get", lambda _: get_package_mock(data))
-    monkeypatch.setattr(pypi_provider, "get_metadata_for_wheel", lambda _: get_metadata_mock())
+    monkeypatch.setattr(
+        pypi_provider.Candidate, "_get_metadata_for_wheel", lambda _: get_metadata_mock()
+    )
 
     resolver = resolvelib.ResolveLibResolver()
     req = Requirement("flask==2.0.1")
@@ -215,7 +223,9 @@ def test_resolvelib_sdist_invalid_suffix(monkeypatch):
     data = '<a href="https://example.com/Flask-2.0.1.foo">Flask-2.0.1.foo</a><br/>'
 
     monkeypatch.setattr(requests, "get", lambda _: get_package_mock(data))
-    monkeypatch.setattr(pypi_provider, "get_metadata_for_wheel", lambda _: get_metadata_mock())
+    monkeypatch.setattr(
+        pypi_provider.Candidate, "_get_metadata_for_wheel", lambda _: get_metadata_mock()
+    )
 
     resolver = resolvelib.ResolveLibResolver()
     req = Requirement("flask==2.0.1")