From fc7e383d4b765dc91379a5f56f89370d6448abbe Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Mon, 4 Nov 2024 19:10:35 +0900
Subject: [PATCH 01/10] improve the whl_annotation handling

---
 CHANGELOG.md                      | 2 ++
 python/private/pypi/extension.bzl | 8 ++------
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f21f9bb3cc..19ceb95318 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -37,6 +37,8 @@ A brief description of the categories of changes:
   by default. Users wishing to keep this argument and to enforce more hermetic
   builds can do so by passing the argument in
   [`pip.parse#extra_pip_args`](https://rules-python.readthedocs.io/en/latest/api/rules_python/python/extensions/pip.html#pip.parse.extra_pip_args)
+* (pip.parse) {attr}`pip.parse.whl_modifications` now normalizes the given whl names
+  and now `pyyaml` and `PyYAML` will both work.
 
 {#v0-0-0-fixed}
 ### Fixed
diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
index fdc76d5750..5278fc5624 100644
--- a/python/private/pypi/extension.bzl
+++ b/python/private/pypi/extension.bzl
@@ -136,7 +136,7 @@ def _create_whl_repos(
     whl_modifications = {}
     if pip_attr.whl_modifications != None:
         for mod, whl_name in pip_attr.whl_modifications.items():
-            whl_modifications[whl_name] = mod
+            whl_modifications[normalize_name(whl_name)] = mod
 
     if pip_attr.experimental_requirement_cycles:
         requirement_cycles = {
@@ -214,10 +214,6 @@ def _create_whl_repos(
 
     repository_platform = host_platform(module_ctx)
     for whl_name, requirements in requirements_by_platform.items():
-        # We are not using the "sanitized name" because the user
-        # would need to guess what name we modified the whl name
-        # to.
-        annotation = whl_modifications.get(whl_name)
         whl_name = normalize_name(whl_name)
 
         group_name = whl_group_mapping.get(whl_name)
@@ -231,7 +227,7 @@ def _create_whl_repos(
         )
         maybe_args = dict(
             # The following values are safe to omit if they have false like values
-            annotation = annotation,
+            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,

From c0087c484e94b2eb9cd31ae2c3f683073edc426b Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Mon, 4 Nov 2024 19:13:03 +0900
Subject: [PATCH 02/10] improve the whl_annotation handling

---
 tests/pypi/extension/extension_tests.bzl | 39 +++++++++++++++++++++---
 1 file changed, 35 insertions(+), 4 deletions(-)

diff --git a/tests/pypi/extension/extension_tests.bzl b/tests/pypi/extension/extension_tests.bzl
index 27c6bba5e2..04a11047c9 100644
--- a/tests/pypi/extension/extension_tests.bzl
+++ b/tests/pypi/extension/extension_tests.bzl
@@ -20,12 +20,12 @@ load("//python/private/pypi:extension.bzl", "parse_modules")  # buildifier: disa
 
 _tests = []
 
-def _mock_mctx(*modules, environ = {}, read = None):
+def _mock_mctx(*modules, environ = {}, read = None, os_name = "unittest", os_arch = "exotic"):
     return struct(
-        os = struct(
+        os = os or struct(
             environ = environ,
-            name = "unittest",
-            arch = "exotic",
+            name = os_name,
+            arch = os_arch,
         ),
         read = read or (lambda _: "simple==0.0.1 --hash=sha256:deadbeef"),
         modules = [
@@ -158,6 +158,37 @@ def _test_simple(env):
 
 _tests.append(_test_simple)
 
+def _test_simple_with_whl_mods(env):
+    pypi = _parse_modules(
+        env,
+        module_ctx = _mock_mctx(
+            _mod(
+                name = "rules_python",
+                parse = [
+                    _parse(
+                        hub_name = "pypi",
+                        python_version = "3.15",
+                        requirements_lock = "requirements.txt",
+                    ),
+                ],
+            ),
+            os_name = "linux",
+            os_arch = "aarch64",
+        ),
+        available_interpreters = {
+            "python_3_15_host": "unit_test_interpreter_target",
+        },
+    )
+
+    pypi.is_reproducible().equals(True)
+    pypi.exposed_packages().contains_exactly({"pypi": []})
+    pypi.hub_group_map().contains_exactly({"pypi": {}})
+    pypi.hub_whl_map().contains_exactly({"pypi": {}})
+    pypi.whl_libraries().contains_exactly({})
+    pypi.whl_mods().contains_exactly({})
+
+_tests.append(_test_simple_with_whl_mods)
+
 def _test_simple_get_index(env):
     got_simpleapi_download_args = []
     got_simpleapi_download_kwargs = {}

From f9532feebf00779ee90abe6a95c3f94c39af81a1 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Mon, 4 Nov 2024 19:15:00 +0900
Subject: [PATCH 03/10] add a test

---
 tests/pypi/extension/extension_tests.bzl | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/tests/pypi/extension/extension_tests.bzl b/tests/pypi/extension/extension_tests.bzl
index 04a11047c9..7fce6ae5d5 100644
--- a/tests/pypi/extension/extension_tests.bzl
+++ b/tests/pypi/extension/extension_tests.bzl
@@ -22,7 +22,7 @@ _tests = []
 
 def _mock_mctx(*modules, environ = {}, read = None, os_name = "unittest", os_arch = "exotic"):
     return struct(
-        os = os or struct(
+        os = struct(
             environ = environ,
             name = os_name,
             arch = os_arch,
@@ -183,8 +183,17 @@ def _test_simple_with_whl_mods(env):
     pypi.is_reproducible().equals(True)
     pypi.exposed_packages().contains_exactly({"pypi": []})
     pypi.hub_group_map().contains_exactly({"pypi": {}})
-    pypi.hub_whl_map().contains_exactly({"pypi": {}})
-    pypi.whl_libraries().contains_exactly({})
+    pypi.hub_whl_map().contains_exactly({"pypi": {
+        "simple": [struct(config_setting = "//_config:is_python_3.15", filename = None, repo = "pypi_315_simple", target_platforms = None, version = "3.15")],
+    }})
+    pypi.whl_libraries().contains_exactly({
+        "pypi_315_simple": {
+            "dep_template": "@pypi//{name}:{target}",
+            "python_interpreter_target": "unit_test_interpreter_target",
+            "repo": "pypi_315",
+            "requirement": "simple==0.0.1 --hash=sha256:deadbeef",
+        },
+    })
     pypi.whl_mods().contains_exactly({})
 
 _tests.append(_test_simple_with_whl_mods)

From bfd33d66c3e15fcd113dcfd4003a844b61d78d98 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Mon, 4 Nov 2024 19:15:16 +0900
Subject: [PATCH 04/10] fixup! add a test

---
 tests/pypi/extension/extension_tests.bzl | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/tests/pypi/extension/extension_tests.bzl b/tests/pypi/extension/extension_tests.bzl
index 7fce6ae5d5..fab8a22d86 100644
--- a/tests/pypi/extension/extension_tests.bzl
+++ b/tests/pypi/extension/extension_tests.bzl
@@ -184,7 +184,15 @@ def _test_simple_with_whl_mods(env):
     pypi.exposed_packages().contains_exactly({"pypi": []})
     pypi.hub_group_map().contains_exactly({"pypi": {}})
     pypi.hub_whl_map().contains_exactly({"pypi": {
-        "simple": [struct(config_setting = "//_config:is_python_3.15", filename = None, repo = "pypi_315_simple", target_platforms = None, version = "3.15")],
+        "simple": [
+            struct(
+                config_setting = "//_config:is_python_3.15",
+                filename = None,
+                repo = "pypi_315_simple",
+                target_platforms = None,
+                version = "3.15",
+            ),
+        ],
     }})
     pypi.whl_libraries().contains_exactly({
         "pypi_315_simple": {

From d7a3c2767fabe29cbe9c61b7a1800e70711ba720 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Mon, 4 Nov 2024 19:50:52 +0900
Subject: [PATCH 05/10] wip: it seems that due to circular dependencies I
 cannot auto-detect the values of the extra aliases field to be passed to the
 hub repo, because I have no visibility into what annotatinos are passed to
 what hubs, because of how the code is structured today. Lets just ask the
 user to specify those explicitly

---
 python/private/pypi/extension.bzl        | 23 ++++++++++
 tests/pypi/extension/extension_tests.bzl | 57 +++++++++++++++++++++++-
 2 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
index 5278fc5624..a39e3fa56a 100644
--- a/python/private/pypi/extension.bzl
+++ b/python/private/pypi/extension.bzl
@@ -68,6 +68,7 @@ def _create_whl_repos(
         pip_attr,
         whl_overrides,
         simpleapi_cache,
+        whl_mods,
         available_interpreters = INTERPRETER_LABELS,
         simpleapi_download = simpleapi_download):
     """create all of the whl repositories
@@ -104,6 +105,7 @@ def _create_whl_repos(
     # containers to aggregate outputs from this function
     whl_map = {}
     exposed_packages = {}
+    extra_aliases = {}
     whl_libraries = {}
 
     # if we do not have the python_interpreter set in the attributes
@@ -138,6 +140,14 @@ def _create_whl_repos(
         for mod, whl_name in pip_attr.whl_modifications.items():
             whl_modifications[normalize_name(whl_name)] = mod
 
+            annotation = json.decode(module_ctx.read(mod))
+            if annotation.additive_build_content:
+                extra_aliases.setdefault(whl_name, {}).update({
+                    line.partition("name = ")[-1].strip("\"',"): True
+                    for line in annotation.build_content.split("\n")
+                    if line.lstrip().beginswith("name = ")
+                })
+
     if pip_attr.experimental_requirement_cycles:
         requirement_cycles = {
             name: [normalize_name(whl_name) for whl_name in whls]
@@ -349,6 +359,7 @@ def _create_whl_repos(
         is_reproducible = is_reproducible,
         whl_map = whl_map,
         exposed_packages = exposed_packages,
+        extra_aliases = extra_aliases,
         whl_libraries = whl_libraries,
     )
 
@@ -433,6 +444,7 @@ You cannot use both the additive_build_content and additive_build_content_file a
     hub_whl_map = {}
     hub_group_map = {}
     exposed_packages = {}
+    extra_aliases = {}
     whl_libraries = {}
 
     is_reproducible = True
@@ -477,11 +489,15 @@ You cannot use both the additive_build_content and additive_build_content_file a
                 pip_attr = pip_attr,
                 simpleapi_cache = simpleapi_cache,
                 whl_overrides = whl_overrides,
+                whl_mods = whl_mods,
                 **kwargs
             )
             hub_whl_map.setdefault(hub_name, {})
             for key, settings in out.whl_map.items():
                 hub_whl_map[hub_name].setdefault(key, []).extend(settings)
+            extra_aliases.setdefault(hub_name, {})
+            for whl_name, extra_aliases in out.extra_aliases.items():
+                extra_aliases[hub_name].setdefault(whl_name, {}).update(extra_aliases)
             exposed_packages.setdefault(hub_name, {}).update(out.exposed_packages)
             whl_libraries.update(out.whl_libraries)
             is_reproducible = is_reproducible and out.is_reproducible
@@ -510,6 +526,13 @@ You cannot use both the additive_build_content and additive_build_content_file a
             for hub_name, group_map in sorted(hub_group_map.items())
         },
         exposed_packages = {k: sorted(v) for k, v in sorted(exposed_packages.items())},
+        extra_aliases = {
+            hub_name: {
+                whl_name: sorted(aliases)
+                for whl_name, aliases in extra_whl_aliases.items()
+            }
+            for hub_name, extra_whl_aliases in extra_aliases.items()
+        },
         whl_libraries = dict(sorted(whl_libraries.items())),
         is_reproducible = is_reproducible,
     )
diff --git a/tests/pypi/extension/extension_tests.bzl b/tests/pypi/extension/extension_tests.bzl
index fab8a22d86..4a5fbc00ec 100644
--- a/tests/pypi/extension/extension_tests.bzl
+++ b/tests/pypi/extension/extension_tests.bzl
@@ -61,6 +61,7 @@ def _parse_modules(env, **kwargs):
         attrs = dict(
             is_reproducible = subjects.bool,
             exposed_packages = subjects.dict,
+            extra_aliases = subjects.dict,
             hub_group_map = subjects.dict,
             hub_whl_map = subjects.dict,
             whl_libraries = subjects.dict,
@@ -68,6 +69,29 @@ def _parse_modules(env, **kwargs):
         ),
     )
 
+def _whl_mods(
+        *,
+        whl_name,
+        hub_name,
+        additive_build_content = None,
+        additive_build_content_file = None,
+        copy_executables = {},
+        copy_files = {},
+        data = [],
+        data_exclude_glob = [],
+        srcs_exclude_glob = []):
+    return struct(
+        additive_build_content = additive_build_content,
+        additive_build_content_file = additive_build_content_file,
+        copy_executables = copy_executables,
+        copy_files = copy_files,
+        data = data,
+        data_exclude_glob = data_exclude_glob,
+        hub_name = hub_name,
+        srcs_exclude_glob = srcs_exclude_glob,
+        whl_name = whl_name,
+    )
+
 def _parse(
         *,
         hub_name,
@@ -164,11 +188,25 @@ def _test_simple_with_whl_mods(env):
         module_ctx = _mock_mctx(
             _mod(
                 name = "rules_python",
+                whl_mods = [
+                    _whl_mods(
+                        additive_build_content = """\
+filegroup(
+    name = "foo",
+    srcs = ["all"],
+)""",
+                        hub_name = "whl_mods_hub",
+                        whl_name = "simple",
+                    ),
+                ],
                 parse = [
                     _parse(
                         hub_name = "pypi",
                         python_version = "3.15",
                         requirements_lock = "requirements.txt",
+                        whl_modifications = {
+                            "@whl_mods_hub//:simple.json": "simple",
+                        },
                     ),
                 ],
             ),
@@ -182,6 +220,11 @@ def _test_simple_with_whl_mods(env):
 
     pypi.is_reproducible().equals(True)
     pypi.exposed_packages().contains_exactly({"pypi": []})
+    pypi.extra_aliases().contains_exactly({
+        "pypi": {
+            "simple": "foo",
+        }
+    })
     pypi.hub_group_map().contains_exactly({"pypi": {}})
     pypi.hub_whl_map().contains_exactly({"pypi": {
         "simple": [
@@ -196,13 +239,25 @@ def _test_simple_with_whl_mods(env):
     }})
     pypi.whl_libraries().contains_exactly({
         "pypi_315_simple": {
+            "annotation": "@whl_mods_hub//:simple.json",
             "dep_template": "@pypi//{name}:{target}",
             "python_interpreter_target": "unit_test_interpreter_target",
             "repo": "pypi_315",
             "requirement": "simple==0.0.1 --hash=sha256:deadbeef",
         },
     })
-    pypi.whl_mods().contains_exactly({})
+    pypi.whl_mods().contains_exactly({
+        "whl_mods_hub": {
+            "simple": struct(
+                build_content = "filegroup(\n    name = \"foo\",\n    srcs = [\"all\"],\n)",
+                copy_executables = {},
+                copy_files = {},
+                data = [],
+                data_exclude_glob = [],
+                srcs_exclude_glob = [],
+            ),
+        },
+    })
 
 _tests.append(_test_simple_with_whl_mods)
 

From 3b9d26e32960e011dea8ba02d24db503b162b87d Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Mon, 4 Nov 2024 20:09:17 +0900
Subject: [PATCH 06/10] fix(pip.parse): allow users to add extra aliases to the
 hub repo.

Fixes #2187
---
 CHANGELOG.md                               |  3 +++
 examples/bzlmod/BUILD.bazel                | 14 +++++++---
 examples/bzlmod/MODULE.bazel               |  3 +++
 python/private/pypi/extension.bzl          | 30 ++++++++++++----------
 python/private/pypi/hub_repository.bzl     |  5 ++++
 python/private/pypi/render_pkg_aliases.bzl | 24 +++++++++++------
 tests/pypi/extension/extension_tests.bzl   |  9 ++++---
 7 files changed, 61 insertions(+), 27 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 19ceb95318..d5b757e02e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -60,6 +60,9 @@ A brief description of the categories of changes:
   and one extra file `requirements_universal.txt` if you prefer a single file.
   The `requirements.txt` file may be removed in the future.
 * The rules_python version is now reported in `//python/features.bzl#features.version`
+* (pip.parse) {attr}`pip.parse.extra_hub_aliases` can now be used to expose extra
+  targets created by annotations in whl repositories.
+  Fixes [#2187](https://github.com/bazelbuild/rules_python/issues/2187).
 
 {#v0-0-0-removed}
 ### Removed
diff --git a/examples/bzlmod/BUILD.bazel b/examples/bzlmod/BUILD.bazel
index d684b9c31d..054b957b3b 100644
--- a/examples/bzlmod/BUILD.bazel
+++ b/examples/bzlmod/BUILD.bazel
@@ -69,16 +69,24 @@ py_test_with_transition(
 # to run some of the tests.
 # See: https://github.com/bazelbuild/bazel-skylib/blob/main/docs/build_test_doc.md
 build_test(
-    name = "all_wheels",
+    name = "all_wheels_build_test",
     targets = all_whl_requirements,
 )
 
 build_test(
-    name = "all_data_requirements",
+    name = "all_data_requirements_build_test",
     targets = all_data_requirements,
 )
 
 build_test(
-    name = "all_requirements",
+    name = "all_requirements_build_test",
     targets = all_requirements,
 )
+
+# Check the annotations API
+build_test(
+    name = "extra_annotation_targets_build_test",
+    targets = [
+        "@pip//wheel:generated_file",
+    ],
+)
diff --git a/examples/bzlmod/MODULE.bazel b/examples/bzlmod/MODULE.bazel
index 57e6b7e7b7..4381cb0d70 100644
--- a/examples/bzlmod/MODULE.bazel
+++ b/examples/bzlmod/MODULE.bazel
@@ -183,6 +183,9 @@ pip.parse(
         "cp39_linux_*",
         "cp39_*",
     ],
+    extra_hub_aliases = {
+        "wheel": ["generated_file"],
+    },
     hub_name = "pip",
     python_version = "3.9",
     requirements_lock = "requirements_lock_3_9.txt",
diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
index a39e3fa56a..c5660272ec 100644
--- a/python/private/pypi/extension.bzl
+++ b/python/private/pypi/extension.bzl
@@ -68,7 +68,6 @@ def _create_whl_repos(
         pip_attr,
         whl_overrides,
         simpleapi_cache,
-        whl_mods,
         available_interpreters = INTERPRETER_LABELS,
         simpleapi_download = simpleapi_download):
     """create all of the whl repositories
@@ -105,7 +104,10 @@ def _create_whl_repos(
     # containers to aggregate outputs from this function
     whl_map = {}
     exposed_packages = {}
-    extra_aliases = {}
+    extra_aliases = {
+        whl_name: {alias: True for alias in aliases}
+        for whl_name, aliases in pip_attr.extra_hub_aliases.items()
+    }
     whl_libraries = {}
 
     # if we do not have the python_interpreter set in the attributes
@@ -140,14 +142,6 @@ def _create_whl_repos(
         for mod, whl_name in pip_attr.whl_modifications.items():
             whl_modifications[normalize_name(whl_name)] = mod
 
-            annotation = json.decode(module_ctx.read(mod))
-            if annotation.additive_build_content:
-                extra_aliases.setdefault(whl_name, {}).update({
-                    line.partition("name = ")[-1].strip("\"',"): True
-                    for line in annotation.build_content.split("\n")
-                    if line.lstrip().beginswith("name = ")
-                })
-
     if pip_attr.experimental_requirement_cycles:
         requirement_cycles = {
             name: [normalize_name(whl_name) for whl_name in whls]
@@ -489,15 +483,14 @@ You cannot use both the additive_build_content and additive_build_content_file a
                 pip_attr = pip_attr,
                 simpleapi_cache = simpleapi_cache,
                 whl_overrides = whl_overrides,
-                whl_mods = whl_mods,
                 **kwargs
             )
             hub_whl_map.setdefault(hub_name, {})
             for key, settings in out.whl_map.items():
                 hub_whl_map[hub_name].setdefault(key, []).extend(settings)
             extra_aliases.setdefault(hub_name, {})
-            for whl_name, extra_aliases in out.extra_aliases.items():
-                extra_aliases[hub_name].setdefault(whl_name, {}).update(extra_aliases)
+            for whl_name, aliases in out.extra_aliases.items():
+                extra_aliases[hub_name].setdefault(whl_name, {}).update(aliases)
             exposed_packages.setdefault(hub_name, {}).update(out.exposed_packages)
             whl_libraries.update(out.whl_libraries)
             is_reproducible = is_reproducible and out.is_reproducible
@@ -617,6 +610,7 @@ def _pip_impl(module_ctx):
         hub_repository(
             name = hub_name,
             repo_name = hub_name,
+            extra_hub_aliases = mods.extra_aliases.get(hub_name, {}),
             whl_map = {
                 key: json.encode(value)
                 for key, value in whl_map.items()
@@ -703,6 +697,16 @@ The indexes must support Simple API as described here:
 https://packaging.python.org/en/latest/specifications/simple-repository-api/
 """,
         ),
+        "extra_hub_aliases": attr.string_list_dict(
+            doc = """\
+Extra aliases to make for specific wheels in the hub repo. This is useful when
+paired with the {attr}`whl_modifications`.
+
+:::{versionadded} 0.38.0
+:::
+""",
+            mandatory = False,
+        ),
         "hub_name": attr.string(
             mandatory = True,
             doc = """
diff --git a/python/private/pypi/hub_repository.bzl b/python/private/pypi/hub_repository.bzl
index 7afb616e3d..69d937142a 100644
--- a/python/private/pypi/hub_repository.bzl
+++ b/python/private/pypi/hub_repository.bzl
@@ -35,6 +35,7 @@ def _impl(rctx):
             key: [whl_alias(**v) for v in json.decode(values)]
             for key, values in rctx.attr.whl_map.items()
         },
+        extra_hub_aliases = rctx.attr.extra_hub_aliases,
         requirement_cycles = rctx.attr.groups,
     )
     for path, contents in aliases.items():
@@ -65,6 +66,10 @@ def _impl(rctx):
 
 hub_repository = repository_rule(
     attrs = {
+        "extra_hub_aliases": attr.string_list_dict(
+            doc = "Extra aliases to make for specific wheels in the hub repo.",
+            mandatory = True,
+        ),
         "groups": attr.string_list_dict(
             mandatory = False,
         ),
diff --git a/python/private/pypi/render_pkg_aliases.bzl b/python/private/pypi/render_pkg_aliases.bzl
index 0086bfff8f..60f4b54306 100644
--- a/python/private/pypi/render_pkg_aliases.bzl
+++ b/python/private/pypi/render_pkg_aliases.bzl
@@ -117,7 +117,7 @@ def _render_whl_library_alias(
         **kwargs
     )
 
-def _render_common_aliases(*, name, aliases, group_name = None):
+def _render_common_aliases(*, name, aliases, extra_aliases = [], group_name = None):
     lines = [
         """load("@bazel_skylib//lib:selects.bzl", "selects")""",
         """package(default_visibility = ["//visibility:public"])""",
@@ -153,12 +153,17 @@ def _render_common_aliases(*, name, aliases, group_name = None):
                 target_name = target_name,
                 visibility = ["//_groups:__subpackages__"] if name.startswith("_") else None,
             )
-            for target_name, name in {
-                PY_LIBRARY_PUBLIC_LABEL: PY_LIBRARY_IMPL_LABEL if group_name else PY_LIBRARY_PUBLIC_LABEL,
-                WHEEL_FILE_PUBLIC_LABEL: WHEEL_FILE_IMPL_LABEL if group_name else WHEEL_FILE_PUBLIC_LABEL,
-                DATA_LABEL: DATA_LABEL,
-                DIST_INFO_LABEL: DIST_INFO_LABEL,
-            }.items()
+            for target_name, name in (
+                {
+                    PY_LIBRARY_PUBLIC_LABEL: PY_LIBRARY_IMPL_LABEL if group_name else PY_LIBRARY_PUBLIC_LABEL,
+                    WHEEL_FILE_PUBLIC_LABEL: WHEEL_FILE_IMPL_LABEL if group_name else WHEEL_FILE_PUBLIC_LABEL,
+                    DATA_LABEL: DATA_LABEL,
+                    DIST_INFO_LABEL: DIST_INFO_LABEL,
+                } | {
+                    x: x
+                    for x in extra_aliases
+                }
+            ).items()
         ],
     )
     if group_name:
@@ -177,7 +182,7 @@ def _render_common_aliases(*, name, aliases, group_name = None):
 
     return "\n\n".join(lines)
 
-def render_pkg_aliases(*, aliases, requirement_cycles = None):
+def render_pkg_aliases(*, aliases, requirement_cycles = None, extra_hub_aliases = {}):
     """Create alias declarations for each PyPI package.
 
     The aliases should be appended to the pip_repository BUILD.bazel file. These aliases
@@ -188,6 +193,8 @@ def render_pkg_aliases(*, aliases, requirement_cycles = None):
         aliases: dict, the keys are normalized distribution names and values are the
             whl_alias instances.
         requirement_cycles: any package groups to also add.
+        extra_hub_aliases: The list of extra aliases for each whl to be added
+          in addition to the default ones.
 
     Returns:
         A dict of file paths and their contents.
@@ -215,6 +222,7 @@ def render_pkg_aliases(*, aliases, requirement_cycles = None):
         "{}/BUILD.bazel".format(normalize_name(name)): _render_common_aliases(
             name = normalize_name(name),
             aliases = pkg_aliases,
+            extra_aliases = extra_hub_aliases.get(name, []),
             group_name = whl_group_mapping.get(normalize_name(name)),
         ).strip()
         for name, pkg_aliases in aliases.items()
diff --git a/tests/pypi/extension/extension_tests.bzl b/tests/pypi/extension/extension_tests.bzl
index 4a5fbc00ec..aa120af83d 100644
--- a/tests/pypi/extension/extension_tests.bzl
+++ b/tests/pypi/extension/extension_tests.bzl
@@ -105,6 +105,7 @@ def _parse(
         experimental_index_url = "",
         experimental_requirement_cycles = {},
         experimental_target_platforms = [],
+        extra_hub_aliases = {},
         extra_pip_args = [],
         isolated = True,
         netrc = None,
@@ -130,6 +131,7 @@ def _parse(
         experimental_index_url = experimental_index_url,
         experimental_requirement_cycles = experimental_requirement_cycles,
         experimental_target_platforms = experimental_target_platforms,
+        extra_hub_aliases = extra_hub_aliases,
         extra_pip_args = extra_pip_args,
         hub_name = hub_name,
         isolated = isolated,
@@ -204,6 +206,9 @@ filegroup(
                         hub_name = "pypi",
                         python_version = "3.15",
                         requirements_lock = "requirements.txt",
+                        extra_hub_aliases = {
+                            "simple": ["foo"],
+                        },
                         whl_modifications = {
                             "@whl_mods_hub//:simple.json": "simple",
                         },
@@ -221,9 +226,7 @@ filegroup(
     pypi.is_reproducible().equals(True)
     pypi.exposed_packages().contains_exactly({"pypi": []})
     pypi.extra_aliases().contains_exactly({
-        "pypi": {
-            "simple": "foo",
-        }
+        "pypi": {"simple": ["foo"]},
     })
     pypi.hub_group_map().contains_exactly({"pypi": {}})
     pypi.hub_whl_map().contains_exactly({"pypi": {

From 743aa21b9aac1e8f94e94f912eb45dad5cfe389f Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Mon, 4 Nov 2024 23:21:23 +0900
Subject: [PATCH 07/10] pre-commit: update lock file

---
 examples/bzlmod/MODULE.bazel.lock | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/examples/bzlmod/MODULE.bazel.lock b/examples/bzlmod/MODULE.bazel.lock
index d34f4ecb73..93cdf9d9ac 100644
--- a/examples/bzlmod/MODULE.bazel.lock
+++ b/examples/bzlmod/MODULE.bazel.lock
@@ -1392,8 +1392,8 @@
     },
     "@@rules_python~//python/extensions:pip.bzl%pip": {
       "general": {
-        "bzlTransitiveDigest": "qxyKk6sb6G2WeW3iUlRmVO5jafUab5qPwz66Y2anPp8=",
-        "usagesDigest": "MChlcSw99EuW3K7OOoMcXQIdcJnEh6YmfyjJm+9mxIg=",
+        "bzlTransitiveDigest": "E5Yr6AjquyIy5ae3c7URmvtPPOm2j+7XOr58GOHp8vw=",
+        "usagesDigest": "lk6zXfe4e1mBpnQyGpa/uFKLzACIpsI82I00SPeiOYs=",
         "recordedFileInputs": {
           "@@other_module~//requirements_lock_3_11.txt": "a7d0061366569043d5efcf80e34a32c732679367cb3c831c4cdc606adc36d314",
           "@@rules_python~//python/private/pypi/whl_installer/platform.py": "b944b908b25a2f97d6d9f491504ad5d2507402d7e37c802ee878783f87f2aa11",
@@ -3239,6 +3239,7 @@
             "ruleClassName": "hub_repository",
             "attributes": {
               "repo_name": "other_module_pip",
+              "extra_hub_aliases": {},
               "whl_map": {
                 "absl_py": "[{\"config_setting\":\"//_config:is_python_3.11\",\"filename\":null,\"repo\":\"other_module_pip_311_absl_py\",\"target_platforms\":null,\"version\":\"3.11\"}]"
               },
@@ -4564,6 +4565,11 @@
             "ruleClassName": "hub_repository",
             "attributes": {
               "repo_name": "pip",
+              "extra_hub_aliases": {
+                "wheel": [
+                  "generated_file"
+                ]
+              },
               "whl_map": {
                 "alabaster": "[{\"config_setting\":\"//_config:is_python_3.10\",\"filename\":null,\"repo\":\"pip_310_alabaster\",\"target_platforms\":null,\"version\":\"3.10\"},{\"config_setting\":\"//_config:is_python_3.9\",\"filename\":\"alabaster-0.7.13-py3-none-any.whl\",\"repo\":\"pip_39_alabaster_py3_none_any_1ee19aca\",\"target_platforms\":null,\"version\":\"3.9\"},{\"config_setting\":\"//_config:is_python_3.9\",\"filename\":\"alabaster-0.7.13.tar.gz\",\"repo\":\"pip_39_alabaster_sdist_a27a4a08\",\"target_platforms\":null,\"version\":\"3.9\"}]",
                 "astroid": "[{\"config_setting\":\"//_config:is_python_3.10\",\"filename\":null,\"repo\":\"pip_310_astroid\",\"target_platforms\":null,\"version\":\"3.10\"},{\"config_setting\":\"//_config:is_python_3.9\",\"filename\":\"astroid-2.12.13-py3-none-any.whl\",\"repo\":\"pip_39_astroid_py3_none_any_10e0ad5f\",\"target_platforms\":null,\"version\":\"3.9\"},{\"config_setting\":\"//_config:is_python_3.9\",\"filename\":\"astroid-2.12.13.tar.gz\",\"repo\":\"pip_39_astroid_sdist_1493fe8b\",\"target_platforms\":null,\"version\":\"3.9\"}]",
@@ -6581,7 +6587,7 @@
     },
     "@@rules_python~//python/private/pypi:pip.bzl%pip_internal": {
       "general": {
-        "bzlTransitiveDigest": "6NoEDGeQugmtzNzf4Emcb8Sb/cW3RTxSSA6DTHLB1/A=",
+        "bzlTransitiveDigest": "wz5L+/+R6gOtD681pNVgPUUipqqPH0bP/b0e22JbSOI=",
         "usagesDigest": "LYtSAPzhPjmfD9vF39mCED1UQSvHEo2Hv+aK5Z4ZWWc=",
         "recordedFileInputs": {
           "@@rules_python~//tools/publish/requirements_linux.txt": "8175b4c8df50ae2f22d1706961884beeb54e7da27bd2447018314a175981997d",
@@ -8765,6 +8771,7 @@
             "ruleClassName": "hub_repository",
             "attributes": {
               "repo_name": "rules_python_publish_deps",
+              "extra_hub_aliases": {},
               "whl_map": {
                 "backports_tarfile": "[{\"config_setting\":\"//_config:is_python_3.11\",\"filename\":\"backports.tarfile-1.2.0-py3-none-any.whl\",\"repo\":\"rules_python_publish_deps_311_backports_tarfile_py3_none_any_77e284d7\",\"target_platforms\":null,\"version\":\"3.11\"},{\"config_setting\":\"//_config:is_python_3.11\",\"filename\":\"backports_tarfile-1.2.0.tar.gz\",\"repo\":\"rules_python_publish_deps_311_backports_tarfile_sdist_d75e02c2\",\"target_platforms\":null,\"version\":\"3.11\"}]",
                 "certifi": "[{\"config_setting\":\"//_config:is_python_3.11\",\"filename\":\"certifi-2024.8.30-py3-none-any.whl\",\"repo\":\"rules_python_publish_deps_311_certifi_py3_none_any_922820b5\",\"target_platforms\":null,\"version\":\"3.11\"},{\"config_setting\":\"//_config:is_python_3.11\",\"filename\":\"certifi-2024.8.30.tar.gz\",\"repo\":\"rules_python_publish_deps_311_certifi_sdist_bec941d2\",\"target_platforms\":null,\"version\":\"3.11\"}]",

From 7e0029b4ff3463255fc9ad12e6df718eab99e37d Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Tue, 5 Nov 2024 08:28:19 +0900
Subject: [PATCH 08/10] bazelci: change the integration test for lockfile

---
 .bazelci/presubmit.yml | 19 ++++++++-----------
 1 file changed, 8 insertions(+), 11 deletions(-)

diff --git a/.bazelci/presubmit.yml b/.bazelci/presubmit.yml
index 5d51b106ed..4da960fb6a 100644
--- a/.bazelci/presubmit.yml
+++ b/.bazelci/presubmit.yml
@@ -65,15 +65,6 @@ buildifier:
 .reusable_build_test_all: &reusable_build_test_all
   build_targets: ["..."]
   test_targets: ["..."]
-.lockfile_mode_error: &lockfile_mode_error
-  # For testing lockfile support
-  skip_in_bazel_downstream_pipeline: "Lockfile depends on the bazel version"
-  build_flags:
-    - "--lockfile_mode=error"
-  test_flags:
-    - "--lockfile_mode=error"
-  coverage_flags:
-    - "--lockfile_mode=error"
 .coverage_targets_example_bzlmod: &coverage_targets_example_bzlmod
   coverage_targets: ["..."]
 .coverage_targets_example_bzlmod_build_file_generation: &coverage_targets_example_bzlmod_build_file_generation
@@ -268,17 +259,23 @@ tasks:
   integration_test_bzlmod_ubuntu_lockfile:
     <<: *reusable_build_test_all
     <<: *coverage_targets_example_bzlmod
-    <<: *lockfile_mode_error
     name: "examples/bzlmod: Ubuntu with lockfile"
     working_directory: examples/bzlmod
     platform: ubuntu2004
+    shell_commands:
+    # Update the lockfiles and fail if it is different.
+    - "./tools/private/update_bzlmod_lockfiles.sh"
+    - "git diff --exit-code"
   integration_test_bzlmod_macos_lockfile:
     <<: *reusable_build_test_all
     <<: *coverage_targets_example_bzlmod
-    <<: *lockfile_mode_error
     name: "examples/bzlmod: macOS with lockfile"
     working_directory: examples/bzlmod
     platform: macos
+    shell_commands:
+    # Update the lockfiles and fail if it is different.
+    - "./tools/private/update_bzlmod_lockfiles.sh"
+    - "git diff --exit-code"
 
   integration_test_bzlmod_generate_build_file_generation_ubuntu_min:
     <<: *minimum_supported_version

From 2fcd3eba82b5f4abd4c9cbb5649c645fba326265 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Tue, 5 Nov 2024 08:28:48 +0900
Subject: [PATCH 09/10] update the lockfile

---
 examples/bzlmod/MODULE.bazel.lock | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/bzlmod/MODULE.bazel.lock b/examples/bzlmod/MODULE.bazel.lock
index 93cdf9d9ac..c115ef974f 100644
--- a/examples/bzlmod/MODULE.bazel.lock
+++ b/examples/bzlmod/MODULE.bazel.lock
@@ -1393,7 +1393,7 @@
     "@@rules_python~//python/extensions:pip.bzl%pip": {
       "general": {
         "bzlTransitiveDigest": "E5Yr6AjquyIy5ae3c7URmvtPPOm2j+7XOr58GOHp8vw=",
-        "usagesDigest": "lk6zXfe4e1mBpnQyGpa/uFKLzACIpsI82I00SPeiOYs=",
+        "usagesDigest": "iVxh/vcpGrSKpO8rafQwAe7uq+pHhasSXC7Pg4o/1dw=",
         "recordedFileInputs": {
           "@@other_module~//requirements_lock_3_11.txt": "a7d0061366569043d5efcf80e34a32c732679367cb3c831c4cdc606adc36d314",
           "@@rules_python~//python/private/pypi/whl_installer/platform.py": "b944b908b25a2f97d6d9f491504ad5d2507402d7e37c802ee878783f87f2aa11",

From 68edd7d6c0f621f2ecc9f1c93070aa670f28f744 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Tue, 5 Nov 2024 08:34:30 +0900
Subject: [PATCH 10/10] fixup! bazelci: change the integration test for
 lockfile

---
 .bazelci/presubmit.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.bazelci/presubmit.yml b/.bazelci/presubmit.yml
index 4da960fb6a..a5f893fa89 100644
--- a/.bazelci/presubmit.yml
+++ b/.bazelci/presubmit.yml
@@ -264,7 +264,7 @@ tasks:
     platform: ubuntu2004
     shell_commands:
     # Update the lockfiles and fail if it is different.
-    - "./tools/private/update_bzlmod_lockfiles.sh"
+    - "../../tools/private/update_bzlmod_lockfiles.sh"
     - "git diff --exit-code"
   integration_test_bzlmod_macos_lockfile:
     <<: *reusable_build_test_all
@@ -274,7 +274,7 @@ tasks:
     platform: macos
     shell_commands:
     # Update the lockfiles and fail if it is different.
-    - "./tools/private/update_bzlmod_lockfiles.sh"
+    - "../../tools/private/update_bzlmod_lockfiles.sh"
     - "git diff --exit-code"
 
   integration_test_bzlmod_generate_build_file_generation_ubuntu_min: