diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl index cafe245eb2..1d8f7b1880 100644 --- a/python/private/pypi/extension.bzl +++ b/python/private/pypi/extension.bzl @@ -32,6 +32,10 @@ load(":pypi_cache.bzl", "pypi_cache") load(":simpleapi_download.bzl", "simpleapi_download") load(":unified_hub_repo.bzl", "unified_hub_repo") load(":whl_library.bzl", "whl_library") +load(":whl_archive.bzl", "whl_archive") +load(":whl_deps_repo.bzl", "whl_deps_repo") +load(":whl_repo_name.bzl", "whl_repo_name") +load(":pep508_requirement.bzl", "requirement") def _whl_mods_impl(whl_mods_dict): """Implementation of the pip.whl_mods tag class. @@ -537,6 +541,149 @@ def _create_unified_hub_repo(mods): packages = packages, ) +def register_whl_libraries( + registrations, + rules = struct( + whl_library=whl_library, + whl_deps_repo=whl_deps_repo, + whl_archive=whl_archive, + ), + ): + """Register all of the whl libraries that where produced by the extension. + + Args: + registrations: {type}`dict[str, dict[str, Any]]` The args to pass to + the repository rules when processing everything. + rules: {type}`struct` used to inject different implementations used to + create instances. Used for testing. + """ + # we first check which wheels are ambiguous - belong to the same hub repository, indicated + # by the `dep_template`, but have different (URL, index_url) pairs for each whl. + urls = {} + fallback = {} + separate_srcs = {} + for name, args in registrations.items(): + if "urls" not in args or "filename" not in args: + # only enabled with bazel downloader + continue + + if "annotation" in args or "whl_patches" in args: + # No support for patched repos yet + continue + + # For each each integrity, filename, we should have a unique url and index_url + # used across all of the hub repos, if not, print a warning and fallback to the + # previous behaviour of not using the source repository. + # + # If the filename is the same but it has a different sha256, then it means that + # there could be one of the following situations: + # 1. dependency confusion. + # 2. somebody patched the wheel. + # 3. it is an unlucky coincidence. + # + # Any of these cases are a bit tricky to get right, the repo name is derived + # filename together. The repo name is derived from the filename and a hash for + # this reason. + integrity = args["integrity"] + filename = args["filename"] + + url = args["urls"][0] + index_url = args["index_url"] + + hub_name, _, _ = args["dep_template"].partition("//") + hub_name = hub_name.strip("@") + + # for now we include the hub_name in the key since the reuse is done for each + # hub_name separately. + key = (hub_name, filename, integrity) + value = (url, index_url) + + if (filename, integrity) in fallback: + continue + + got = urls.setdefault(key, value) + if got != value: + urls.pop(key) + fallback.setdefault((filename, integrity), None) + + separate_srcs[name] = key + + for name, args in registrations.items(): + if name in separate_srcs: + # first handle the regular case + continue + + rules.whl_library(name = name, **args) + + # Then handle the reusable case + registered_srcs = {} + for name, key in separate_srcs.items(): + args = registrations[name] + hub_name, _, integrity = key + + filename = args["filename"] + extract_args = { + k: v + for k, v in args.items() + if k not in { + "config_load": None, + "dep_template": None, + } + } + + # The extras do not affect the extraction, so normalize the requirement + # to allow the same wheel to be extracted only once. + extract_args["requirement"] = _without_extras(extract_args["requirement"]) + extract_repo_name = "{}_{}".format( + hub_name, + whl_repo_name(filename, ""), + ) + if key not in registered_srcs: + registered_srcs[key] = extract_repo_name + rules.whl_archive( + name = extract_repo_name, + **extract_args, + ) + + deps_args = { + k: args.get(k) + for k in [ + "config_load", + "requirement", + "dep_template", + "group_deps", + "group_name", + "repo_prefix", + ] + if args.get(k) != None + } | { + "metadata_file": "@{}//:metadata.json".format(extract_repo_name), + } + rules.whl_deps_repo( + name = name, + **deps_args + ) + +def _without_extras(requirement_line): + """Remove the extras from a requirement line. + + The extras do not affect which wheel is downloaded and extracted, so they + can be removed to allow the same wheel to be extracted only once even when + it is referenced with different extras. + + Args: + requirement_line: {type}`str` the requirement line, e.g. + `foo[bar]==1.0 --hash=sha256:...`. + + Returns: + The requirement line without the extras, e.g. `foo==1.0 --hash=sha256:...`. + """ + name_and_version, _, extras = requirement_line.partition("[") + if not extras: + return requirement_line + _, _, rest = extras.partition("]") + return name_and_version + rest + def _pip_impl(module_ctx): """Implementation of a class tag that creates the pip hub and corresponding pip spoke whl repositories. @@ -611,8 +758,7 @@ def _pip_impl(module_ctx): # Build all of the wheel modifications if the tag class is called. _whl_mods_impl(mods.whl_mods) - for name, args in mods.whl_libraries.items(): - whl_library(name = name, **args) + register_whl_libraries(registrations = mods.whl_libraries) for hub_name, whl_map in mods.hub_whl_map.items(): hub_repository( diff --git a/python/private/pypi/whl_library_deps_targets.bzl b/python/private/pypi/whl_library_deps_targets.bzl index 89e9b74dc7..a9b63b7322 100644 --- a/python/private/pypi/whl_library_deps_targets.bzl +++ b/python/private/pypi/whl_library_deps_targets.bzl @@ -23,8 +23,8 @@ def whl_library_deps_targets( repo, aliases = None, metadata_name, - requires_dist, - extras, + requires_dist = [], + extras = [], include = [], group_deps = [], group_name = None, @@ -168,6 +168,7 @@ def whl_library_deps_targets( package_deps = package_deps, tmpl = dep_template.format(name = "{}", target = PY_LIBRARY_PUBLIC_LABEL), ), + precompile = "disabled", tags = tags, visibility = impl_vis, ) diff --git a/python/private/pypi/whl_library_targets.bzl b/python/private/pypi/whl_library_targets.bzl index b8d51d0bd7..4f7bf8e98a 100644 --- a/python/private/pypi/whl_library_targets.bzl +++ b/python/private/pypi/whl_library_targets.bzl @@ -86,16 +86,17 @@ def whl_library_targets( **kwargs ) - whl_library_deps_targets( - name = name, - metadata_name = metadata_name, - requires_dist = requires_dist, - dep_template = dep_template, # only needed if requires_dist or group_name is present - group_deps = group_deps, # only needed if group_name is present - group_name = group_name, # must specify group_deps together - extras = extras, # only needed if requires_dist is present - include = include, # only needed if requires_dist is present - repo = None, # set aliases in the same repo - aliases = {}, - **kwargs - ) + if dep_template: + whl_library_deps_targets( + name = name, + metadata_name = metadata_name, + requires_dist = requires_dist, + dep_template = dep_template, # only needed if requires_dist or group_name is present + group_deps = group_deps, # only needed if group_name is present + group_name = group_name, # must specify group_deps together + extras = extras, # only needed if requires_dist is present + include = include, # only needed if requires_dist is present + repo = None, # set aliases in the same repo + aliases = {}, + **kwargs + )