feat: allow overriding rules in whl_library repository rule

CHANGELOG.md

@@ -31,7 +31,13 @@ A brief description of the categories of changes:
 * (gazelle): Fix incorrect use of `t.Fatal`/`t.Fatalf` in tests.
 
 ### Added
-* Nothing yet
+* (pypi): Add macro wrappers for all of the publicly exposed targets in the
+  `whl_library` repository_rule. This allows users to override rules used to
+  extract whl targets from the `.whl` distributions. This allows `WORKSPACE`
+  and `bzlmod` users to override the load statements. The `WORKSPACE` can pass
+  an extra argument named `override_loads` to the `pip_install` macro from the
+  hub repository, whereas `bzlmod` users have a new attribute `load_symbols` in
+  the `pip.override` tag class. TODO: link to the ticket and add documentation.
 
 ### Removed
 * Nothing yet

examples/bzlmod/MODULE.bazel

@@ -195,6 +195,12 @@ pip.parse(
 # The patches have to be in the unified-diff format.
 pip.override(
     file = "requests-2.25.1-py2.py3-none-any.whl",
+    # One can also override the loads for better support of using alternative implementations
+    # For allowed values, inspect the error message when typing any unsupported
+    # value as the key of the dictionary.
+    library_symbols = {
+        "py_binary": "@rules_python//python:defs.bzl",
+    },
     patch_strip = 1,
     patches = [
         "@//patches:empty.patch",

python/private/pypi/extension.bzl

@@ -238,6 +238,9 @@ def _create_whl_repos(module_ctx, pip_attr, whl_map, whl_overrides, group_map, s
             repo = pip_name,
             dep_template = "@{}//{{name}}:{{target}}".format(hub_name),
         )
+
+        overrides = whl_overrides.get(whl_name)
+
         maybe_args = dict(
             # The following values are safe to omit if they have false like values
             annotation = annotation,
@@ -251,10 +254,11 @@ def _create_whl_repos(module_ctx, pip_attr, whl_map, whl_overrides, group_map, s
             pip_data_exclude = pip_attr.pip_data_exclude,
             python_interpreter = pip_attr.python_interpreter,
             python_interpreter_target = python_interpreter_target,
+            override_loads = overrides.loads if overrides else None,
             whl_patches = {
                 p: json.encode(args)
-                for p, args in whl_overrides.get(whl_name, {}).items()
-            },
+                for p, args in overrides.patches.items()
+            } if overrides else None,
         )
         whl_library_args.update({k: v for k, v in maybe_args.items() if v})
         maybe_args_with_default = dict(
@@ -441,17 +445,19 @@ def _pip_impl(module_ctx):
                 fail("Duplicate module overrides for '{}'".format(attr.file))
             _overriden_whl_set[attr.file] = None
 
-            for patch in attr.patches:
-                if whl_name not in whl_overrides:
-                    whl_overrides[whl_name] = {}
+            overrides = whl_overrides.setdefault(whl_name, struct(
+                patches = {},
+                loads = attr.library_symbols or {},
+            ))
 
-                if patch not in whl_overrides[whl_name]:
-                    whl_overrides[whl_name][patch] = struct(
+            for patch in attr.patches:
+                overrides.patches.setdefault(
+                    patch,
+                    struct(
                         patch_strip = attr.patch_strip,
                         whls = [],
-                    )
-
-                whl_overrides[whl_name][patch].whls.append(attr.file)
+                    ),
+                ).whls.append(attr.file)
 
     # Used to track all the different pip hubs and the spoke pip Python
     # versions.
@@ -727,6 +733,15 @@ applied to all repositories that setup this distribution via the pip.parse tag
 class.""",
             mandatory = True,
         ),
+        "library_symbols": attr.string_dict(
+            doc = """
+The string dictionary for symbols to be used when defining targets within the `whl_library`.
+
+This allows users to override the rules used for particular wheels for better
+support of generating `py_library` from an `sdist` or potentially improve how
+the `whl_filegroup` defines providers.
+""",
+        ),
         "patch_strip": attr.int(
             default = 0,
             doc = """\

python/private/pypi/generate_whl_library_build_bazel.bzl

@@ -27,6 +27,9 @@ load(
     "WHEEL_FILE_PUBLIC_LABEL",
 )
 
+_DEFAULT_MACRO_LOAD = "@rules_python//python/private/pypi:whl_library_macros.bzl"
+_COPY_FILE_LOAD = "@bazel_skylib//rules:copy_file.bzl"
+
 _COPY_FILE_TEMPLATE = """\
 copy_file(
     name = "{dest}.copy",
@@ -52,24 +55,15 @@ _BUILD_TEMPLATE = """\
 
 package(default_visibility = ["//visibility:public"])
 
-filegroup(
-    name = "{dist_info_label}",
-    srcs = glob(["site-packages/*.dist-info/**"], allow_empty = True),
-)
-
-filegroup(
-    name = "{data_label}",
-    srcs = glob(["data/**"], allow_empty = True),
-)
-
-filegroup(
+data_filegroup(name="{data_label}")
+dist_info_filegroup(name="{dist_info_label}")
+whl_file(
     name = "{whl_file_label}",
     srcs = ["{whl_name}"],
-    data = {whl_file_deps},
+    deps = {whl_file_deps},
     visibility = {impl_vis},
 )
-
-py_library(
+whl_library(
     name = "{py_library_label}",
     srcs = glob(
         ["site-packages/**/*.py"],
@@ -78,14 +72,11 @@ py_library(
         # pure-Python code, e.g. pymssql, which is written in Cython.
         allow_empty = True,
     ),
+    deps = {dependencies},
     data = {data} + glob(
         ["site-packages/**/*"],
         exclude={data_exclude},
     ),
-    # This makes this directory a top-level in the python import
-    # search path for anything that depends on this.
-    imports = ["site-packages"],
-    deps = {dependencies},
     tags = {tags},
     visibility = {impl_vis},
 )
@@ -126,7 +117,7 @@ def _render_list_and_select(deps, deps_by_platform, tmpl):
         return "{} + {}".format(deps, deps_by_platform)
 
 def _render_config_settings(dependencies_by_platform):
-    loads = []
+    loads = {}
     additional_content = []
     for p in dependencies_by_platform:
         # p can be one of the following formats:
@@ -158,7 +149,7 @@ def _render_config_settings(dependencies_by_platform):
 
         if abi:
             if not loads:
-                loads.append("""load("@rules_python//python/config_settings:config_settings.bzl", "is_python_config_setting")""")
+                loads["is_python_config_setting"] = "@rules_python//python/config_settings:config_settings.bzl"
 
             additional_content.append(
                 """\
@@ -197,6 +188,7 @@ def generate_whl_library_build_bazel(
         data_exclude,
         tags,
         entry_points,
+        override_loads = {},
         annotation = None,
         group_name = None,
         group_deps = []):
@@ -217,6 +209,9 @@ def generate_whl_library_build_bazel(
         group_deps: List[str]; names of fellow members of the group (if any). These will be excluded
           from generated deps lists so as to avoid direct cycles. These dependencies will be provided
           at runtime by the group rules which wrap this library and its fellows together.
+        override_loads: dict[str, str], the dictionary for the symbols to be
+            used for defining standard targets. If the key within dict does not
+            correspond to a symbol, it will fail.
 
     Returns:
         A complete BUILD file as a string
@@ -290,16 +285,25 @@ def generate_whl_library_build_bazel(
         if deps
     }
 
-    loads = [
-        """load("@rules_python//python:defs.bzl", "py_library", "py_binary")""",
-        """load("@bazel_skylib//rules:copy_file.bzl", "copy_file")""",
-    ]
+    loads = {
+        "copy_file": _COPY_FILE_LOAD,
+        "data_filegroup": _DEFAULT_MACRO_LOAD,
+        "dist_info_filegroup": _DEFAULT_MACRO_LOAD,
+        "py_binary": "@rules_python//python:py_binary.bzl",
+        "whl_file": _DEFAULT_MACRO_LOAD,
+        "whl_library": _DEFAULT_MACRO_LOAD,
+    }
+    for symbol, location in override_loads.items():
+        if symbol in loads:
+            loads[symbol] = location
+        else:
+            msg = "Unsupported symbol name '{}', use one of: {}".format(symbol, sorted(loads))
+            fail(msg)
 
     loads_, config_settings_content = _render_config_settings(dependencies_by_platform)
     if config_settings_content:
-        for line in loads_:
-            if line not in loads:
-                loads.append(line)
+        for symbol, loc in loads_.items():
+            loads[symbol] = loc
         additional_content.append(config_settings_content)
 
     lib_dependencies = _render_list_and_select(
@@ -356,7 +360,7 @@ def generate_whl_library_build_bazel(
     contents = "\n".join(
         [
             _BUILD_TEMPLATE.format(
-                loads = "\n".join(sorted(loads)),
+                loads = _render_loads(loads),
                 py_library_label = py_library_label,
                 dependencies = render.indent(lib_dependencies, " " * 4).lstrip(),
                 whl_file_deps = render.indent(whl_file_deps, " " * 4).lstrip(),
@@ -366,7 +370,6 @@ def generate_whl_library_build_bazel(
                 tags = repr(tags),
                 data_label = DATA_LABEL,
                 dist_info_label = DIST_INFO_LABEL,
-                entry_point_prefix = WHEEL_ENTRY_POINT_PREFIX,
                 srcs_exclude = repr(srcs_exclude),
                 data = repr(data),
                 impl_vis = repr(impl_vis),
@@ -418,3 +421,21 @@ def _generate_entry_point_rule(*, name, script, pkg):
         src = script.replace("\\", "/"),
         pkg = pkg,
     )
+
+def _render_loads(loads):
+    by_import = {}
+    for symbol, loc in loads.items():
+        by_import.setdefault(loc, []).append(symbol)
+
+    lines = []
+    for loc, symbols in sorted(by_import.items()):
+        if len(symbols) == 1:
+            line = "load({}, {})".format(repr(loc), repr(symbols[0]))
+        else:
+            line = "load(\n{}\n)".format(render.indent("\n".join(
+                [repr(item) + "," for item in [loc] + sorted(symbols)],
+            )))
+
+        lines.append(line)
+
+    return "\n".join(lines)

python/private/pypi/whl_library.bzl

@@ -333,6 +333,7 @@ def _whl_library_impl(rctx):
         group_name = rctx.attr.group_name,
         group_deps = rctx.attr.group_deps,
         data_exclude = rctx.attr.pip_data_exclude,
+        override_loads = rctx.attr.override_loads,
         tags = [
             "pypi_name=" + metadata["name"],
             "pypi_version=" + metadata["version"],
@@ -398,6 +399,15 @@ and the target that we need respectively.
     "group_name": attr.string(
         doc = "Name of the group, if any.",
     ),
+    "override_loads": attr.string_dict(
+        doc = """
+The string dictionary for symbols to be used when defining targets within the `whl_library`.
+
+This allows users to override the rules used for particular wheels for better
+support of generating `py_library` from an `sdist` or potentially improve how
+the `whl_filegroup` defines providers.
+""",
+    ),
     "repo": attr.string(
         mandatory = True,
         doc = "Pointer to parent repo name. Used to make these rules rerun if the parent repo changes.",

python/private/pypi/whl_library_macros.bzl

@@ -0,0 +1,114 @@
+# Copyright 2024 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Generate targets for the whl_library macro."""
+
+load("//python:py_library.bzl", "py_library")
+load(":labels.bzl", "DATA_LABEL", "DIST_INFO_LABEL")
+
+def dist_info_filegroup(
+        *,
+        name = DIST_INFO_LABEL,
+        visibility = ["//visibility:public"],
+        native = native):
+    """Generate the dist-info target.
+
+    Args:
+        name: The name for the `dist_info` target.
+        visibility: The visibility of the target.
+        native: The native struct for unit testing.
+    """
+    native.filegroup(
+        name = name,
+        srcs = native.glob(["site-packages/*.dist-info/**"], allow_empty = True),
+        visibility = visibility,
+    )
+
+def data_filegroup(
+        *,
+        name = DATA_LABEL,
+        visibility = ["//visibility:public"],
+        native = native):
+    """Generate the data target.
+
+    Args:
+        name: The name for the `data` target.
+        visibility: The visibility of the target.
+        native: The native struct for unit testing.
+    """
+    native.filegroup(
+        name = name,
+        srcs = native.glob(["data/**"], allow_empty = True),
+        visibility = visibility,
+    )
+
+def whl_file(
+        *,
+        name,
+        deps,
+        srcs,
+        visibility = [],
+        native = native):
+    """Generate the whl target.
+
+    Args:
+        name: None, unused.
+        deps: The whl deps.
+        srcs: The list of whl sources.
+        visibility: The visibility passed to the whl target in order
+            to group dependencies.
+        native: The native struct for unit testing.
+    """
+    native.filegroup(
+        name = name,
+        srcs = srcs,
+        data = deps,
+        visibility = visibility,
+    )
+
+def whl_library(
+        *,
+        name,
+        data,
+        deps,
+        srcs,
+        tags = [],
+        visibility = [],
+        py_library = py_library,
+        native = native):
+    """Generate the targets that are exposed by an extracted whl library.
+
+    Args:
+        name: the name of the library target.
+        data: The py_library data.
+        deps: The py_library dependencies.
+        srcs: The python srcs.
+        tags: The tags set to the py_library target to force rebuilding when
+            the version of the dependencies changes.
+        visibility: The visibility passed to the whl and py_library targets in order
+            to group dependencies.
+        py_library: The py_library rule to use for defining the targets.
+        native: The native struct for unit testing.
+    """
+    py_library(
+        name = name,
+        srcs = srcs,
+        data = data,
+        # This makes this directory a top-level in the python import
+        # search path for anything that depends on this.
+        imports = ["site-packages"],
+        deps = deps,
+        tags = tags,
+        visibility = visibility,
+    )