fix accidentally double-registering setuptools for ipex (pantsbuild#9341

)

When creating `.ipex` files as per pantsbuild#8793, we don't add resolved distributions to `self._distributions` in `PexBuilderWrapper`. This leads to the annoying `_sanitize_requirements()` hack introduced into `ipex_launcher.py` in that PR (having to remove duplicate `setuptools` requirements).

- Add all resolved distributions to `self._distributions` in `PexBuilderWrapper`, even when `generate_ipex=True`.
- Remove the separate list of transitively-resolved requirements passed through several methods in `PexBuilderWrapper`.

src/python/pants/backend/python/subsystems/ipex/ipex_launcher.py

@@ -17,7 +17,6 @@
 from pex.interpreter import PythonInterpreter
 from pex.pex_builder import PEXBuilder
 from pex.pex_info import PexInfo
-from pkg_resources import Requirement
 
 
 APP_CODE_PREFIX = 'user_files/'
@@ -34,26 +33,6 @@ def _log(message):
   sys.stderr.write(message + '\n')
 
 
-def _sanitize_requirements(requirements):
-  """
-  Remove duplicate keys such as setuptools or pex which may be injected multiple times into the
-  resulting ipex when first executed.
-  """
-  project_names = []
-  new_requirements = {}
-
-  for r in requirements:
-    r = Requirement(r)
-    if r.marker and not r.marker.evaluate():
-      continue
-    if r.name not in new_requirements:
-      project_names.append(r.name)
-      new_requirements[r.name] = str(r)
-  sanitized_requirements = [new_requirements[n] for n in project_names]
-
-  return sanitized_requirements
-
-
 def modify_pex_info(pex_info, **kwargs):
   new_info = json.loads(pex_info.dump())
   new_info.update(kwargs)
@@ -88,10 +67,6 @@ def _hydrate_pex_file(self, hydrated_pex_file):
   fetchers.extend(Fetcher([url]) for url in resolver_settings.pop('find_links'))
   resolver_settings['fetchers'] = fetchers
 
-  sanitized_requirements = _sanitize_requirements(bootstrap_info.requirements)
-  bootstrap_info = modify_pex_info(bootstrap_info, requirements=sanitized_requirements)
-  bootstrap_builder.info = bootstrap_info
-
   resolved_distributions = resolver.resolve(
     requirements=bootstrap_info.requirements,
     cache=bootstrap_info.pex_root,

src/python/pants/python/pex_build_util.py

@@ -215,17 +215,13 @@ def extract_single_dist_for_current_platform(self, reqs, dist_key) -> Distributi
     :raises: :class:`self.SingleDistExtractionError` if no dists or multiple dists matched the
         given `dist_key`.
     """
-    distributions, _transitive_requirements = self.resolve_distributions(
-      reqs, platforms=["current"]
-    )
+    distributions = self.resolve_distributions(reqs, platforms=["current"])
     try:
       matched_dist = assert_single_element(
-        list(
           dist
-          for _, dists in distributions.items()
+          for dists in distributions.values()
           for dist in dists
           if dist.key == dist_key
-        )
       )
     except (StopIteration, ValueError) as e:
       raise self.SingleDistExtractionError(
@@ -235,7 +231,7 @@ def extract_single_dist_for_current_platform(self, reqs, dist_key) -> Distributi
 
   def resolve_distributions(
     self, reqs: List[PythonRequirement], platforms: Optional[List[Platform]] = None,
-  ) -> Tuple[Dict[str, List[Distribution]], List[PythonRequirement]]:
+  ) -> Dict[str, List[Distribution]]:
     """Multi-platform dependency resolution.
 
     :param reqs: A list of :class:`PythonRequirement` to resolve.
@@ -255,10 +251,8 @@ def resolve_distributions(
         find_links.add(req.repository)
 
     # Resolve the requirements into distributions.
-    distributions, transitive_requirements = self._resolve_multi(
-      self._builder.interpreter, list(deduped_reqs), platforms, list(find_links),
-    )
-    return (distributions, transitive_requirements)
+    distributions = self._resolve_multi(self._builder.interpreter, list(deduped_reqs), platforms, list(find_links))
+    return distributions
 
   def add_resolved_requirements(
     self,
@@ -280,9 +274,7 @@ def add_resolved_requirements(
                                     pex dependency to the output ipex file, and therefore needs to
                                     override the default behavior of this method.
     """
-    distributions, transitive_requirements = self.resolve_distributions(
-      reqs, platforms=platforms
-    )
+    distributions = self.resolve_distributions(reqs, platforms=platforms)
     locations: Set[str] = set()
     for platform, dists in distributions.items():
       for dist in dists:
@@ -291,25 +283,21 @@ def add_resolved_requirements(
             self._log.debug(
               f"  *AVOIDING* dumping distribution into ipex: .../{os.path.basename(dist.location)}"
             )
+            self._register_distribution(dist)
           else:
             self._log.debug(
               f"  Dumping distribution: .../{os.path.basename(dist.location)}")
 
             self.add_distribution(dist)
         locations.add(dist.location)
-    # In addition to the top-level requirements, we add all the requirements matching the resolved
-    # distributions to the resulting pex. If `generate_ipex=True` is set, we need to have all the
-    # transitive requirements resolved in order to hydrate the .ipex with an intransitive resolve.
-    if self._generate_ipex and not override_ipex_build_do_actually_add_distribution:
-      self.add_direct_requirements(transitive_requirements)
 
   def _resolve_multi(
     self,
     interpreter: PythonInterpreter,
     requirements: List[PythonRequirement],
     platforms: Optional[List[Platform]],
     find_links: Optional[List[str]],
-  ) -> Tuple[Dict[str, List[Distribution]], List[PythonRequirement]]:
+  ) -> Dict[str, List[Distribution]]:
     """Multi-platform dependency resolution for PEX files.
 
     Returns a tuple containing a list of distributions that must be included in order to satisfy a
@@ -337,7 +325,6 @@ def _resolve_multi(
     self._all_find_links.update(OrderedSet(find_links))
 
     distributions: Dict[str, List[Distribution]] = defaultdict(list)
-    transitive_requirements: List[PythonRequirement] = []
 
     fetchers = [
         PyPIFetcher(index_url) for index_url in python_repos.indexes
@@ -363,11 +350,9 @@ def _resolve_multi(
                                             default=True),
       )
       for resolved_dist in resolved_dists:
-        dist = resolved_dist.distribution
-        transitive_requirements.append(dist.as_requirement())
-        distributions[platform].append(dist)
+        distributions[platform].append(resolved_dist.distribution)
 
-    return (distributions, transitive_requirements)
+    return distributions
 
   def _create_source_dumper(self, tgt: Target) -> Callable[[str], None]:
     buildroot = get_buildroot()
@@ -457,6 +442,11 @@ def _shuffle_underlying_pex_builder(self) -> Tuple[PexInfo, Path]:
       self._builder.info,
       self._builder.interpreter.identity)
 
+    # Remove all the original top-level requirements in favor of the transitive == requirements.
+    self._builder.info = ipex_launcher.modify_pex_info(self._builder.info, requirements=[])
+    transitive_reqs = [dist.as_requirement() for dist in self._distributions.values()]
+    self.add_direct_requirements(transitive_reqs)
+
     orig_info = self._builder.info.copy()
 
     orig_chroot = self._builder.chroot()
@@ -467,6 +457,8 @@ def _shuffle_underlying_pex_builder(self) -> Tuple[PexInfo, Path]:
       self._builder.info,
       self._builder.interpreter.identity)
 
+    self._distributions = {}
+
     return (orig_info, Path(orig_chroot.path()))
 
   def _shuffle_original_build_info_into_ipex(self):