diff --git a/python/pip_install/extract_wheels/lib/arguments.py b/python/pip_install/extract_wheels/lib/arguments.py index 46d08a8eb5..63d2afb804 100644 --- a/python/pip_install/extract_wheels/lib/arguments.py +++ b/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 - diff --git a/python/pip_install/parse_requirements_to_bzl/__init__.py b/python/pip_install/parse_requirements_to_bzl/__init__.py index f27e2a2a76..3f08c1a198 100644 --- a/python/pip_install/parse_requirements_to_bzl/__init__.py +++ b/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,12 +75,13 @@ 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] ) @@ -71,13 +89,14 @@ def generate_parsed_requirements_contents(all_args: argparse.Namespace) -> str: [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 " + + "---", + ) arguments.parse_common_args(parser) args = parser.parse_args() diff --git a/python/pip_install/parse_requirements_to_bzl/extract_single_wheel/__init__.py b/python/pip_install/parse_requirements_to_bzl/extract_single_wheel/__init__.py index a46ea2ed24..c2f33bcb73 100644 --- a/python/pip_install/parse_requirements_to_bzl/extract_single_wheel/__init__.py +++ b/python/pip_install/parse_requirements_to_bzl/extract_single_wheel/__init__.py @@ -21,6 +21,10 @@ 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 ---", + ) arguments.parse_common_args(parser) args = parser.parse_args() deserialized_args = dict(vars(args)) @@ -28,10 +32,20 @@ def main() -> None: 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: diff --git a/python/pip_install/pip_repository.bzl b/python/pip_install/pip_repository.bzl index d7d11137a0..9108d183ba 100644 --- a/python/pip_install/pip_repository.bzl +++ b/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 ---" + """, + ), "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 ---", + ), "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.""", +)