Add support for multi-platform pip download to pip_parse

python/pip_install/extract_wheels/lib/arguments.py

@@ -35,11 +35,10 @@ def deserialize_structured_args(args):
     Args:
         args: dict of parsed command line arguments
     """
-    structured_args = ("extra_pip_args", "pip_data_exclude", "environment")
+    structured_args = ("extra_pip_args", "pip_data_exclude", "environment", "pip_platform_definitions")
     for arg_name in structured_args:
         if args.get(arg_name) is not None:
             args[arg_name] = json.loads(args[arg_name])["arg"]
         else:
             args[arg_name] = []
     return args
-

python/pip_install/parse_requirements_to_bzl/__init__.py

@@ -3,7 +3,7 @@
 import textwrap
 import sys
 import shlex
-from typing import List, Tuple
+from typing import Dict, List, Optional, Tuple
 
 from python.pip_install.extract_wheels.lib import bazel, arguments
 from pip._internal.req import parse_requirements, constructors
@@ -36,14 +36,31 @@ def parse_install_requirements(requirements_lock: str, extra_pip_args: List[str]
     return install_req_and_lines
 
 
-def repo_names_and_requirements(install_reqs: List[Tuple[InstallRequirement, str]], repo_prefix: str) -> List[Tuple[str, str]]:
-    return [
-        (
-            bazel.sanitise_name(ir.name, prefix=repo_prefix),
-            line,
-        )
-        for ir, line in install_reqs
-    ]
+class NamesAndRequirements:
+    def __init__(self, whls: List[Tuple[str, str, Optional[str]]], aliases: List[Tuple[str, Dict[str, str]]]):
+        self.whls = whls
+        self.aliases = aliases
+
+
+def repo_names_and_requirements(
+        install_reqs: List[Tuple[InstallRequirement, str]],
+        repo_prefix: str,
+        platforms: Dict[str, str]) -> NamesAndRequirements:
+    whls = []
+    aliases = []
+    for ir, line in install_reqs:
+        generic_name = bazel.sanitise_name(ir.name, prefix=repo_prefix)
+        if not platforms:
+            whls.append((generic_name, line, None))
+        else:
+            select_items = {}
+            for key, platform in platforms.items():
+                prefix = bazel.sanitise_name(platform, prefix=repo_prefix) + "__"
+                name = bazel.sanitise_name(ir.name, prefix=prefix)
+                whls.append((name, line, platform))
+                select_items[key] = "@{name}//:pkg".format(name=name)
+            aliases.append((generic_name, select_items))
+    return NamesAndRequirements(whls, aliases)
 
 
 def generate_parsed_requirements_contents(all_args: argparse.Namespace) -> str:
@@ -58,26 +75,28 @@ def generate_parsed_requirements_contents(all_args: argparse.Namespace) -> str:
     args = dict(vars(all_args))
     args = arguments.deserialize_structured_args(args)
     args.setdefault("python_interpreter", sys.executable)
-    # Pop this off because it wont be used as a config argument to the whl_library rule.
+    # Pop these off because they won't be used as a config argument to the whl_library rule.
     requirements_lock = args.pop("requirements_lock")
+    pip_platform_definitions = args.pop("pip_platform_definitions")
     repo_prefix = bazel.whl_library_repo_prefix(args["repo"])
 
     install_req_and_lines = parse_install_requirements(requirements_lock, args["extra_pip_args"])
-    repo_names_and_reqs = repo_names_and_requirements(install_req_and_lines, repo_prefix)
+    repo_names_and_reqs = repo_names_and_requirements(install_req_and_lines, repo_prefix, pip_platform_definitions)
     all_requirements = ", ".join(
         [bazel.sanitised_repo_library_label(ir.name, repo_prefix=repo_prefix) for ir, _ in install_req_and_lines]
     )
     all_whl_requirements = ", ".join(
         [bazel.sanitised_repo_file_label(ir.name, repo_prefix=repo_prefix) for ir, _ in install_req_and_lines]
     )
     return textwrap.dedent("""\
-        load("@rules_python//python/pip_install:pip_repository.bzl", "whl_library")
+        load("@rules_python//python/pip_install:pip_repository.bzl", "whl_library", "platform_alias")
 
         all_requirements = [{all_requirements}]
 
         all_whl_requirements = [{all_whl_requirements}]
 
-        _packages = {repo_names_and_reqs}
+        _packages = {whl_definitions}
+        _aliases = {alias_definitions}
         _config = {args}
 
         def _clean_name(name):
@@ -90,18 +109,26 @@ def whl_requirement(name):
            return "@{repo_prefix}" + _clean_name(name) + "//:whl"
 
         def install_deps():
-            for name, requirement in _packages:
+            for name, requirement, platform in _packages:
                 whl_library(
                     name = name,
                     requirement = requirement,
+                    pip_platform_definition = platform,
                     **_config,
                 )
+            for name, select_items in _aliases:
+                platform_alias(
+                    name = name,
+                    select_items = select_items,
+                )
         """.format(
             all_requirements=all_requirements,
             all_whl_requirements=all_whl_requirements,
-            repo_names_and_reqs=repo_names_and_reqs,
+            whl_definitions=repo_names_and_reqs.whls,
+            alias_definitions=repo_names_and_reqs.aliases,
             args=args,
             repo_prefix=repo_prefix,
+            pip_platform_definitions=pip_platform_definitions,
             )
         )
 
@@ -133,6 +160,11 @@ def main() -> None:
         required=True,
         help="timeout to use for pip operation.",
     )
+    parser.add_argument(
+        "--pip_platform_definitions",
+        help="A map of select keys to platform definitions in the form "
+             + "<platform>-<python_version>-<implementation>-<abi>",
+    )
     arguments.parse_common_args(parser)
     args = parser.parse_args()
 

python/pip_install/parse_requirements_to_bzl/extract_single_wheel/__init__.py

@@ -21,17 +21,31 @@ def main() -> None:
         required=True,
         help="A single PEP508 requirement specifier string.",
     )
+    parser.add_argument(
+        "--pip_platform_definition",
+        help="A pip platform definition in the form <platform>-<python_version>-<implementation>-<abi>",
+    )
     arguments.parse_common_args(parser)
     args = parser.parse_args()
     deserialized_args = dict(vars(args))
     arguments.deserialize_structured_args(deserialized_args)
 
     configure_reproducible_wheels()
 
-    pip_args = (
-        [sys.executable, "-m", "pip", "--isolated", "wheel", "--no-deps"] +
-        deserialized_args["extra_pip_args"]
-    )
+    pip_args = [sys.executable, "-m", "pip", "--isolated"]
+    if args.pip_platform_definition:
+        platform, python_version, implementation, abi = args.pip_platform_definition.split("-")
+        pip_args.extend([
+            "download",
+            "--only-binary", ":all:",
+            "--platform", platform,
+            "--python-version", python_version,
+            "--implementation", implementation,
+            "--abi", abi
+        ])
+    else:
+        pip_args.append("wheel")
+    pip_args.extend(["--no-deps"] + deserialized_args["extra_pip_args"])
 
     requirement_file = NamedTemporaryFile(mode='wb', delete=False)
     try:

python/pip_install/pip_repository.bzl

@@ -100,6 +100,11 @@ def _pip_repository_impl(rctx):
             "--timeout",
             str(rctx.attr.timeout),
         ]
+        if rctx.attr.pip_platform_definitions:
+            args.extend([
+                "--pip_platform_definitions",
+                struct(arg = {str(k): v for k, v in rctx.attr.pip_platform_definitions.items()}).to_json(),
+            ])
     else:
         args = [
             python_interpreter,
@@ -178,6 +183,11 @@ pip_repository_attrs = {
         default = False,
         doc = "Create the repository in incremental mode.",
     ),
+    "pip_platform_definitions": attr.label_keyed_string_dict(
+        doc = """
+A map of select keys to platform definitions in the form <platform>-<python_version>-<implementation>-<abi>"
+        """,
+    ),
     "requirements": attr.label(
         allow_single_file = True,
         doc = "A 'requirements.txt' pip requirements file.",
@@ -252,6 +262,12 @@ def _impl_whl_library(rctx):
     ]
     args = _parse_optional_attrs(rctx, args)
 
+    if rctx.attr.pip_platform_definition:
+        args.extend([
+            "--pip_platform_definition",
+            rctx.attr.pip_platform_definition,
+        ])
+
     result = rctx.execute(
         args,
         # Manually construct the PYTHONPATH since we cannot use the toolchain here
@@ -266,6 +282,9 @@ def _impl_whl_library(rctx):
     return
 
 whl_library_attrs = {
+    "pip_platform_definition": attr.string(
+        doc = "A pip platform definition in the form <platform>-<python_version>-<implementation>-<abi>",
+    ),
     "repo": attr.string(
         mandatory = True,
         doc = "Pointer to parent repo name. Used to make these rules rerun if the parent repo changes.",
@@ -285,3 +304,29 @@ Download and extracts a single wheel based into a bazel repo based on the requir
 Instantiated from pip_repository and inherits config options from there.""",
     implementation = _impl_whl_library,
 )
+
+_PLATFORM_ALIAS_TMPL = """
+alias(
+    name = "pkg",
+    actual = select({select_items}),
+    visibility = ["//visibility:public"],
+)
+"""
+
+def _impl_platform_alias(rctx):
+    rctx.file(
+        "BUILD",
+        content = _PLATFORM_ALIAS_TMPL.format(
+            select_items = rctx.attr.select_items,
+        ),
+        executable = False,
+    )
+
+platform_alias = repository_rule(
+    attrs = {
+        "select_items": attr.string_dict(),
+    },
+    implementation = _impl_platform_alias,
+    doc = """
+An internal rule used to create an alias for a pip package for the appropriate platform.""",
+)