refactor(pypi): return a list from parse_requirements (#2931)

The modeling of the data structures returned by the `parse_requirements`
function was not optimal and this was because historically there was
more logic in the `extension.bzl` and more things were decided there.

With the recent refactors it is possible to have a harder to misuse data
structure from the `parse_requirements`. For each `package` we will
return a struct which will have a `srcs` field that will contain easy to
consume values.

With this in place we can do the fix that is outlined in the referenced
issue.

Work towards #2648

Author: aignas

python/private/pypi/extension.bzl

@@ -202,8 +202,12 @@ def _create_whl_repos(
         logger = logger,
     )
 
-    for whl_name, requirements in requirements_by_platform.items():
-        group_name = whl_group_mapping.get(whl_name)
+    exposed_packages = {}
+    for whl in requirements_by_platform:
+        if whl.is_exposed:
+            exposed_packages[whl.name] = None
+
+        group_name = whl_group_mapping.get(whl.name)
         group_deps = requirement_cycles.get(group_name, [])
 
         # Construct args separately so that the lock file can be smaller and does not include unused
@@ -214,7 +218,7 @@ def _create_whl_repos(
         maybe_args = dict(
             # The following values are safe to omit if they have false like values
             add_libdir_to_library_search_path = pip_attr.add_libdir_to_library_search_path,
-            annotation = whl_modifications.get(whl_name),
+            annotation = whl_modifications.get(whl.name),
             download_only = pip_attr.download_only,
             enable_implicit_namespace_pkgs = pip_attr.enable_implicit_namespace_pkgs,
             environment = pip_attr.environment,
@@ -226,7 +230,7 @@ def _create_whl_repos(
             python_interpreter_target = python_interpreter_target,
             whl_patches = {
                 p: json.encode(args)
-                for p, args in whl_overrides.get(whl_name, {}).items()
+                for p, args in whl_overrides.get(whl.name, {}).items()
             },
         )
         if not enable_pipstar:
@@ -245,119 +249,99 @@ def _create_whl_repos(
             if v != default
         })
 
-        for requirement in requirements:
-            for repo_name, (args, config_setting) in _whl_repos(
-                requirement = requirement,
+        for src in whl.srcs:
+            repo = _whl_repo(
+                src = src,
                 whl_library_args = whl_library_args,
                 download_only = pip_attr.download_only,
                 netrc = pip_attr.netrc,
                 auth_patterns = pip_attr.auth_patterns,
                 python_version = major_minor,
-                multiple_requirements_for_whl = len(requirements) > 1.,
+                is_multiple_versions = whl.is_multiple_versions,
                 enable_pipstar = enable_pipstar,
-            ).items():
-                repo_name = "{}_{}".format(pip_name, repo_name)
-                if repo_name in whl_libraries:
-                    fail("Attempting to creating a duplicate library {} for {}".format(
-                        repo_name,
-                        whl_name,
-                    ))
+            )
 
-                whl_libraries[repo_name] = args
-                whl_map.setdefault(whl_name, {})[config_setting] = repo_name
+            repo_name = "{}_{}".format(pip_name, repo.repo_name)
+            if repo_name in whl_libraries:
+                fail("Attempting to creating a duplicate library {} for {}".format(
+                    repo_name,
+                    whl.name,
+                ))
+
+            whl_libraries[repo_name] = repo.args
+            whl_map.setdefault(whl.name, {})[repo.config_setting] = repo_name
 
     return struct(
         whl_map = whl_map,
-        exposed_packages = {
-            whl_name: None
-            for whl_name, requirements in requirements_by_platform.items()
-            if len([r for r in requirements if r.is_exposed]) > 0
-        },
+        exposed_packages = exposed_packages,
         extra_aliases = extra_aliases,
         whl_libraries = whl_libraries,
     )
 
-def _whl_repos(*, requirement, whl_library_args, download_only, netrc, auth_patterns, multiple_requirements_for_whl = False, python_version, enable_pipstar = False):
-    ret = {}
-
-    dists = requirement.whls
-    if not download_only and requirement.sdist:
-        dists = dists + [requirement.sdist]
-
-    for distribution in dists:
-        args = dict(whl_library_args)
-        if netrc:
-            args["netrc"] = netrc
-        if auth_patterns:
-            args["auth_patterns"] = auth_patterns
-
-        if not distribution.filename.endswith(".whl"):
-            # pip is not used to download wheels and the python
-            # `whl_library` helpers are only extracting things, however
-            # for sdists, they will be built by `pip`, so we still
-            # need to pass the extra args there.
-            args["extra_pip_args"] = requirement.extra_pip_args
-
-        # This is no-op because pip is not used to download the wheel.
-        args.pop("download_only", None)
-
-        args["requirement"] = requirement.line
-        args["urls"] = [distribution.url]
-        args["sha256"] = distribution.sha256
-        args["filename"] = distribution.filename
-        if not enable_pipstar:
-            args["experimental_target_platforms"] = [
-                # Get rid of the version fot the target platforms because we are
-                # passing the interpreter any way. Ideally we should search of ways
-                # how to pass the target platforms through the hub repo.
-                p.partition("_")[2]
-                for p in requirement.target_platforms
-            ]
-
-        # Pure python wheels or sdists may need to have a platform here
-        target_platforms = None
-        if distribution.filename.endswith(".whl") and not distribution.filename.endswith("-any.whl"):
-            pass
-        elif multiple_requirements_for_whl:
-            target_platforms = requirement.target_platforms
-
-        repo_name = whl_repo_name(
-            distribution.filename,
-            distribution.sha256,
-        )
-        ret[repo_name] = (
-            args,
-            whl_config_setting(
+def _whl_repo(*, src, whl_library_args, is_multiple_versions, download_only, netrc, auth_patterns, python_version, enable_pipstar = False):
+    args = dict(whl_library_args)
+    args["requirement"] = src.requirement_line
+    is_whl = src.filename.endswith(".whl")
+
+    if src.extra_pip_args and not is_whl:
+        # pip is not used to download wheels and the python
+        # `whl_library` helpers are only extracting things, however
+        # for sdists, they will be built by `pip`, so we still
+        # need to pass the extra args there, so only pop this for whls
+        args["extra_pip_args"] = src.extra_pip_args
+
+    if not src.url or (not is_whl and download_only):
+        # Fallback to a pip-installed wheel
+        target_platforms = src.target_platforms if is_multiple_versions else []
+        return struct(
+            repo_name = pypi_repo_name(
+                normalize_name(src.distribution),
+                *target_platforms
+            ),
+            args = args,
+            config_setting = whl_config_setting(
                 version = python_version,
-                filename = distribution.filename,
-                target_platforms = target_platforms,
+                target_platforms = target_platforms or None,
             ),
         )
 
-    if ret:
-        return ret
-
-    # Fallback to a pip-installed wheel
-    args = dict(whl_library_args)  # make a copy
-    args["requirement"] = requirement.line
-    if requirement.extra_pip_args:
-        args["extra_pip_args"] = requirement.extra_pip_args
+    # This is no-op because pip is not used to download the wheel.
+    args.pop("download_only", None)
+
+    if netrc:
+        args["netrc"] = netrc
+    if auth_patterns:
+        args["auth_patterns"] = auth_patterns
+
+    args["urls"] = [src.url]
+    args["sha256"] = src.sha256
+    args["filename"] = src.filename
+    if not enable_pipstar:
+        args["experimental_target_platforms"] = [
+            # Get rid of the version fot the target platforms because we are
+            # passing the interpreter any way. Ideally we should search of ways
+            # how to pass the target platforms through the hub repo.
+            p.partition("_")[2]
+            for p in src.target_platforms
+        ]
+
+    # Pure python wheels or sdists may need to have a platform here
+    target_platforms = None
+    if is_whl and not src.filename.endswith("-any.whl"):
+        pass
+    elif is_multiple_versions:
+        target_platforms = src.target_platforms
 
-    target_platforms = requirement.target_platforms if multiple_requirements_for_whl else []
-    repo_name = pypi_repo_name(
-        normalize_name(requirement.distribution),
-        *target_platforms
-    )
-    ret[repo_name] = (
-        args,
-        whl_config_setting(
+    return struct(
+        repo_name = whl_repo_name(src.filename, src.sha256),
+        args = args,
+        config_setting = whl_config_setting(
             version = python_version,
-            target_platforms = target_platforms or None,
+            filename = src.filename,
+            target_platforms = target_platforms,
         ),
     )
 
-    return ret
-
 def parse_modules(
         module_ctx,
         _fail = fail,

python/private/pypi/parse_requirements.bzl

@@ -179,49 +179,91 @@ def parse_requirements(
             }),
         )
 
-    ret = {}
-    for whl_name, reqs in sorted(requirements_by_platform.items()):
+    ret = []
+    for name, reqs in sorted(requirements_by_platform.items()):
         requirement_target_platforms = {}
         for r in reqs.values():
             target_platforms = env_marker_target_platforms.get(r.requirement_line, r.target_platforms)
             for p in target_platforms:
                 requirement_target_platforms[p] = None
 
-        is_exposed = len(requirement_target_platforms) == len(requirements)
-        if not is_exposed and logger:
+        item = struct(
+            # Return normalized names
+            name = normalize_name(name),
+            is_exposed = len(requirement_target_platforms) == len(requirements),
+            is_multiple_versions = len(reqs.values()) > 1,
+            srcs = _package_srcs(
+                name = name,
+                reqs = reqs,
+                index_urls = index_urls,
+                env_marker_target_platforms = env_marker_target_platforms,
+                extract_url_srcs = extract_url_srcs,
+                logger = logger,
+            ),
+        )
+        ret.append(item)
+        if not item.is_exposed and logger:
             logger.debug(lambda: "Package '{}' will not be exposed because it is only present on a subset of platforms: {} out of {}".format(
-                whl_name,
+                name,
                 sorted(requirement_target_platforms),
                 sorted(requirements),
             ))
 
-        # Return normalized names
-        ret_requirements = ret.setdefault(normalize_name(whl_name), [])
+    if logger:
+        logger.debug(lambda: "Will configure whl repos: {}".format([w.name for w in ret]))
 
-        for r in sorted(reqs.values(), key = lambda r: r.requirement_line):
-            whls, sdist = _add_dists(
-                requirement = r,
-                index_urls = index_urls.get(whl_name),
-                logger = logger,
-            )
+    return ret
 
-            target_platforms = env_marker_target_platforms.get(r.requirement_line, r.target_platforms)
-            ret_requirements.append(
+def _package_srcs(
+        *,
+        name,
+        reqs,
+        index_urls,
+        logger,
+        env_marker_target_platforms,
+        extract_url_srcs):
+    """A function to return sources for a particular package."""
+    srcs = []
+    for r in sorted(reqs.values(), key = lambda r: r.requirement_line):
+        whls, sdist = _add_dists(
+            requirement = r,
+            index_urls = index_urls.get(name),
+            logger = logger,
+        )
+
+        target_platforms = env_marker_target_platforms.get(r.requirement_line, r.target_platforms)
+        target_platforms = sorted(target_platforms)
+
+        all_dists = [] + whls
+        if sdist:
+            all_dists.append(sdist)
+
+        if extract_url_srcs and all_dists:
+            req_line = r.srcs.requirement
+        else:
+            all_dists = [struct(
+                url = "",
+                filename = "",
+                sha256 = "",
+                yanked = False,
+            )]
+            req_line = r.srcs.requirement_line
+
+        for dist in all_dists:
+            srcs.append(
                 struct(
-                    distribution = r.distribution,
-                    line = r.srcs.requirement if extract_url_srcs and (whls or sdist) else r.srcs.requirement_line,
-                    target_platforms = sorted(target_platforms),
+                    distribution = name,
                     extra_pip_args = r.extra_pip_args,
-                    whls = whls,
-                    sdist = sdist,
-                    is_exposed = is_exposed,
+                    requirement_line = req_line,
+                    target_platforms = target_platforms,
+                    filename = dist.filename,
+                    sha256 = dist.sha256,
+                    url = dist.url,
+                    yanked = dist.yanked,
                 ),
             )
 
-    if logger:
-        logger.debug(lambda: "Will configure whl repos: {}".format(ret.keys()))
-
-    return ret
+    return srcs
 
 def select_requirement(requirements, *, platform):
     """A simple function to get a requirement for a particular platform.

python/private/pypi/pip_repository.bzl

@@ -94,15 +94,15 @@ def _pip_repository_impl(rctx):
     selected_requirements = {}
     options = None
     repository_platform = host_platform(rctx)
-    for name, requirements in requirements_by_platform.items():
+    for whl in requirements_by_platform:
         requirement = select_requirement(
-            requirements,
+            whl.srcs,
             platform = None if rctx.attr.download_only else repository_platform,
         )
         if not requirement:
             continue
         options = options or requirement.extra_pip_args
-        selected_requirements[name] = requirement.line
+        selected_requirements[whl.name] = requirement.requirement_line
 
     bzl_packages = sorted(selected_requirements.keys())