Update update-sanity-requirements.py script (#81424)

Frozen requirements can now preserve any explicitly installed package that would normally be omitted, not just setuptools.

1 files changed, 6 insertions, 7 deletions

diff --git a/hacking/update-sanity-requirements.py b/hacking/update-sanity-requirements.py
index 5861590bea..997d6dbf87 100755
--- a/hacking/update-sanity-requirements.py
+++ b/hacking/update-sanity-requirements.py
@@ -15,6 +15,7 @@ import venv
 
 import packaging.version
 import packaging.specifiers
+import packaging.requirements
 
 try:
     import argcomplete
@@ -34,6 +35,11 @@ class SanityTest:
     source_path: pathlib.Path
 
     def freeze_requirements(self) -> None:
+        source_requirements = [packaging.requirements.Requirement(re.sub(' #.*$', '', line)) for line in self.source_path.read_text().splitlines()]
+
+        install_packages = {requirement.name for requirement in source_requirements}
+        exclude_packages = {'distribute', 'pip', 'setuptools', 'wheel'} - install_packages
+
         with tempfile.TemporaryDirectory() as venv_dir:
             venv.create(venv_dir, with_pip=True)
 
@@ -49,13 +55,6 @@ class SanityTest:
             subprocess.run(pip + ['install', 'wheel'], env=env, check=True)  # make bdist_wheel available during pip install
             subprocess.run(pip + ['install', '-r', self.source_path], env=env, check=True)
 
-            keep_setuptools = any(line.startswith('setuptools ') for line in self.source_path.read_text().splitlines())
-
-            exclude_packages = ['pip', 'distribute', 'wheel']
-
-            if not keep_setuptools:
-                exclude_packages.append('setuptools')
-
             freeze_options = ['--all']
 
             for exclude_package in exclude_packages: