From ace01f7e67d38e7c85032bb93d184d739eea46ec Mon Sep 17 00:00:00 2001
From: Yun Peng <pcloudy@google.com>
Date: Tue, 12 May 2026 13:24:02 +0000
Subject: [PATCH 01/11] Add local wheel support

---
 python/private/pypi/extension.bzl            |   4 +
 python/private/pypi/hub_builder.bzl          |  79 +++++++-
 tests/pypi/hub_builder/hub_builder_tests.bzl | 191 +++++++++++++++++++
 tests/support/mocks/mocks.bzl                |  60 +++++-
 tests/support/mocks/mocks_tests.bzl          |   4 +-
 5 files changed, 326 insertions(+), 12 deletions(-)

diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
index 2a8ace28f2..3280786933 100644
--- a/python/private/pypi/extension.bzl
+++ b/python/private/pypi/extension.bzl
@@ -381,6 +381,7 @@ You cannot use both the additive_build_content and additive_build_content_file a
             builder.pip_parse(
                 module_ctx,
                 pip_attr = pip_attr,
+                is_root = mod.is_root,
             )
 
     # Keeps track of all the hub's whl repos across the different versions.
@@ -827,6 +828,9 @@ A dict of labels to wheel names that is typically generated by the whl_modificat
 The labels are JSON config files describing the modifications.
 """,
         ),
+        "local_wheels": attr.string_dict(
+            doc = "Dictionary mapping package names to local wheel file paths relative to the workspace root.",
+        ),
     }, **ATTRS)
     attrs.update(AUTH_ATTRS)
 
diff --git a/python/private/pypi/hub_builder.bzl b/python/private/pypi/hub_builder.bzl
index 1bce648dce..c0a7233a49 100644
--- a/python/private/pypi/hub_builder.bzl
+++ b/python/private/pypi/hub_builder.bzl
@@ -147,7 +147,7 @@ def _build(self):
         whl_libraries = self._whl_libraries,
     )
 
-def _pip_parse(self, module_ctx, pip_attr):
+def _pip_parse(self, module_ctx, pip_attr, is_root = False):
     python_version = pip_attr.python_version
     if python_version in self._platforms:
         fail((
@@ -194,6 +194,7 @@ def _pip_parse(self, module_ctx, pip_attr):
         self,
         module_ctx,
         pip_attr = pip_attr,
+        is_root = is_root,
         enable_pipstar_extract = bool(self._config.enable_pipstar_extract or self._get_index_urls.get(pip_attr.python_version)),
     )
 
@@ -332,7 +333,7 @@ def _add_whl_library(self, *, python_version, whl, repo):
                     if value
                 ])
             ))
-            return
+        return
     self._whl_libraries[repo_name] = repo.args
 
     mapping = self._whl_map.setdefault(whl.name, {})
@@ -479,6 +480,7 @@ def _create_whl_repos(
         module_ctx,
         *,
         pip_attr,
+        is_root = False,
         enable_pipstar_extract = False):
     """create all of the whl repositories
 
@@ -532,7 +534,10 @@ def _create_whl_repos(
 
     interpreter = _detect_interpreter(self, pip_attr)
 
+    local_wheels = _collect_local_wheels(module_ctx, pip_attr, is_root = is_root)
+
     for whl in requirements_by_platform:
+        local_wheel = local_wheels.get(whl.name)
         whl_library_args = common_args | _whl_library_args(
             self,
             whl = whl,
@@ -550,6 +555,7 @@ def _create_whl_repos(
                 python_version = _major_minor_version(pip_attr.python_version),
                 is_multiple_versions = whl.is_multiple_versions,
                 interpreter = interpreter,
+                local_wheel = local_wheel,
                 enable_pipstar_extract = enable_pipstar_extract,
             )
             _add_whl_library(
@@ -625,6 +631,7 @@ def _whl_repo(
         python_version,
         use_downloader,
         interpreter,
+        local_wheel = None,
         enable_pipstar_extract = False):
     args = dict(whl_library_args)
     args["requirement"] = src.requirement_line
@@ -681,9 +688,17 @@ def _whl_repo(
     # targets to each hub for each extra combination and solve this more cleanly as opposed to
     # duplicating whl_library repositories.
     target_platforms = src.target_platforms if is_multiple_versions else []
+    repo_name = whl_repo_name(src.filename, src.sha256, *target_platforms)
+
+    if local_wheel:
+        repo_name += "_local_override"
+        path_str = local_wheel._path if hasattr(local_wheel, "_path") else str(local_wheel)
+        args["urls"] = ["file://" + path_str]
+        args["filename"] = local_wheel.basename
+        args["sha256"] = ""
 
     return struct(
-        repo_name = whl_repo_name(src.filename, src.sha256, *target_platforms),
+        repo_name = repo_name,
         args = args,
         config_setting = whl_config_setting(
             version = python_version,
@@ -696,3 +711,61 @@ def _use_downloader(self, python_version, whl_name):
         normalize_name(whl_name),
         self._get_index_urls.get(python_version) != None,
     )
+
+def _collect_local_wheels(module_ctx, pip_attr, is_root = False):
+    if not is_root:
+        return {}
+
+    wheels = {}
+    explicit_wheels = getattr(pip_attr, "local_wheels", None)
+    if not explicit_wheels:
+        return wheels
+
+    workspace_root = module_ctx.path(Label("@@//:MODULE.bazel")).dirname
+
+    for pkg_name, wheel_path_str in explicit_wheels.items():
+        norm_name = normalize_name(pkg_name)
+        if "*" not in wheel_path_str:
+            wheel_path = workspace_root.get_child(wheel_path_str)
+            if wheel_path.exists:
+                wheels[norm_name] = wheel_path
+        else:
+            last_slash = wheel_path_str.rfind("/")
+            if last_slash >= 0:
+                dir_part = wheel_path_str[:last_slash]
+                pattern = wheel_path_str[last_slash + 1:]
+            else:
+                dir_part = ""
+                pattern = wheel_path_str
+
+            matched_wheel = None
+            target_dir = workspace_root.get_child(dir_part) if dir_part else workspace_root
+            if target_dir.exists:
+                candidates = target_dir.readdir()
+            else:
+                candidates = []
+
+            for candidate in candidates:
+                if not candidate.basename.endswith(".whl"):
+                    continue
+                if _wildcard_match(candidate.basename, pattern):
+                    if not matched_wheel or matched_wheel.basename < candidate.basename:
+                        matched_wheel = candidate
+
+            if matched_wheel:
+                wheels[norm_name] = matched_wheel
+
+    return wheels
+
+def _wildcard_match(name, pattern):
+    if pattern.startswith("*") and pattern.endswith("*"):
+        return name.find(pattern[1:-1]) >= 0
+    elif pattern.startswith("*"):
+        return name.endswith(pattern[1:])
+    elif pattern.endswith("*"):
+        return name.startswith(pattern[:-1])
+    elif "*" in pattern:
+        parts = pattern.split("*", 1)
+        return name.startswith(parts[0]) and name.endswith(parts[1]) and len(name) >= len(parts[0]) + len(parts[1])
+    else:
+        return name == pattern
diff --git a/tests/pypi/hub_builder/hub_builder_tests.bzl b/tests/pypi/hub_builder/hub_builder_tests.bzl
index 216528fc9b..72b46d778b 100644
--- a/tests/pypi/hub_builder/hub_builder_tests.bzl
+++ b/tests/pypi/hub_builder/hub_builder_tests.bzl
@@ -828,6 +828,110 @@ simple==0.0.1 --hash=sha256:deadb00f
 
 _tests.append(_test_index_url_precedence)
 
+def _test_local_wheel_override(env):
+    def mock_simpleapi_download(*_, **__):
+        return {
+            "simple": struct(
+                whls = {
+                    "deadbeef": struct(
+                        yanked = None,
+                        filename = "simple-0.0.1-py3-none-any.whl",
+                        sha256 = "deadbeef",
+                        url = "example.com/simple-0.0.1.whl",
+                    ),
+                },
+                sdists = {},
+                sha256s_by_version = {},
+                index_url = "https://example.com",
+            ),
+        }
+
+    builder = hub_builder(
+        env,
+        simpleapi_download_fn = mock_simpleapi_download,
+    )
+    builder.pip_parse(
+        mocks.mctx(
+            mock_files = {
+                "requirements.txt": "simple==0.0.1 --hash=sha256:deadbeef",
+                "MODULE.bazel": "",
+                "dist/simple-0.0.2-cp315-cp315-linux_x86_64.whl": "",
+            },
+            os_name = "linux",
+            arch_name = "x86_64",
+        ),
+        _parse(
+            hub_name = "pypi",
+            python_version = "3.15",
+            experimental_index_url = "https://example.com",
+            requirements_lock = "requirements.txt",
+            local_wheels = {
+                "simple": "dist/simple-0.0.2-cp315-cp315-linux_x86_64.whl",
+            },
+        ),
+        is_root = True,
+    )
+    pypi = builder.build()
+
+    pypi.exposed_packages().contains_exactly(["simple"])
+    pypi.whl_map().contains_exactly({
+        "simple": {
+            "pypi_315_simple_py3_none_any_deadbeef_local_override": [
+                whl_config_setting(version = "3.15", target_platforms = ["cp315_linux_x86_64"]),
+            ],
+        },
+    })
+    pypi.whl_libraries().contains_exactly({
+        "pypi_315_simple_py3_none_any_deadbeef_local_override": {
+            "config_load": "@pypi//:config.bzl",
+            "dep_template": "@pypi//{name}:{target}",
+            "filename": "simple-0.0.2-cp315-cp315-linux_x86_64.whl",
+            "index_url": "https://example.com",
+            "requirement": "simple==0.0.1",
+            "sha256": "",
+            "urls": ["file://dist/simple-0.0.2-cp315-cp315-linux_x86_64.whl"],
+        },
+    })
+    pypi.extra_aliases().contains_exactly({})
+
+_tests.append(_test_local_wheel_override)
+
+def _test_local_wheel_override_ignored_if_not_root(env):
+    builder = hub_builder(env)
+    builder.pip_parse(
+        mocks.mctx(
+            mock_files = {
+                "requirements.txt": "simple==0.0.1 --hash=sha256:deadbeef",
+                "MODULE.bazel": "",
+                "dist/simple-0.0.2-cp315-cp315-linux_x86_64.whl": "",
+            },
+            os_name = "linux",
+            arch_name = "x86_64",
+        ),
+        _parse(
+            hub_name = "pypi",
+            python_version = "3.15",
+            requirements_lock = "requirements.txt",
+            local_wheels = {
+                "simple": "dist/simple-0.0.2-cp315-cp315-linux_x86_64.whl",
+            },
+        ),
+        is_root = False,
+    )
+    pypi = builder.build()
+
+    pypi.exposed_packages().contains_exactly(["simple"])
+    pypi.whl_libraries().contains_exactly({
+        "pypi_315_simple": {
+            "config_load": "@pypi//:config.bzl",
+            "dep_template": "@pypi//{name}:{target}",
+            "python_interpreter_target": "unit_test_interpreter_target",
+            "requirement": "simple==0.0.1 --hash=sha256:deadbeef",
+        },
+    })
+
+_tests.append(_test_local_wheel_override_ignored_if_not_root)
+
 def _test_download_only_multiple(env):
     builder = hub_builder(env)
     builder.pip_parse(
@@ -1495,6 +1599,93 @@ Attempting to create a duplicate library pypi_315_foo for foo with different arg
 
 _tests.append(_test_err_duplicate_repos)
 
+def _test_explicit_local_wheels(env):
+    def mock_simpleapi_download(*_, **__):
+        return {
+            "simple": struct(
+                whls = {
+                    "deadbeef": struct(
+                        yanked = None,
+                        filename = "simple-0.0.1-py3-none-any.whl",
+                        sha256 = "deadbeef",
+                        url = "example.com/simple-0.0.1.whl",
+                    ),
+                },
+                sdists = {},
+                sha256s_by_version = {},
+                index_url = "https://example.com",
+            ),
+            "libtpu": struct(
+                whls = {
+                    "deadbaaf": struct(
+                        yanked = None,
+                        filename = "libtpu-0.0.1-py3-none-any.whl",
+                        sha256 = "deadbaaf",
+                        url = "example.com/libtpu-0.0.1.whl",
+                    ),
+                },
+                sdists = {},
+                sha256s_by_version = {},
+                index_url = "https://example.com",
+            ),
+        }
+
+    builder = hub_builder(
+        env,
+        simpleapi_download_fn = mock_simpleapi_download,
+    )
+    builder.pip_parse(
+        mocks.mctx(
+            mock_files = {
+                "requirements.txt": """\
+simple==0.0.1 --hash=sha256:deadbeef
+libtpu==0.0.1 --hash=sha256:deadbaaf
+""",
+                "MODULE.bazel": "",
+                "custom_folder/libtpu-0.0.41.dev20260509+nightly-cp314-cp314t-manylinux_2_31_x86_64.whl": "",
+                "custom_folder/simple-0.0.3-py3-none-any.whl": "",
+            },
+            os_name = "linux",
+            arch_name = "x86_64",
+        ),
+        _parse(
+            hub_name = "pypi",
+            python_version = "3.15",
+            experimental_index_url = "https://example.com",
+            requirements_lock = "requirements.txt",
+            local_wheels = {
+                "simple": "custom_folder/simple-0.0.3-py3-none-any.whl",
+                "libtpu": "custom_folder/libtpu-*.whl",
+            },
+        ),
+        is_root = True,
+    )
+    pypi = builder.build()
+
+    pypi.exposed_packages().contains_exactly(["libtpu", "simple"])
+    pypi.whl_libraries().contains_exactly({
+        "pypi_315_libtpu_py3_none_any_deadbaaf_local_override": {
+            "config_load": "@pypi//:config.bzl",
+            "dep_template": "@pypi//{name}:{target}",
+            "filename": "libtpu-0.0.41.dev20260509+nightly-cp314-cp314t-manylinux_2_31_x86_64.whl",
+            "index_url": "https://example.com",
+            "requirement": "libtpu==0.0.1",
+            "sha256": "",
+            "urls": ["file://custom_folder/libtpu-0.0.41.dev20260509+nightly-cp314-cp314t-manylinux_2_31_x86_64.whl"],
+        },
+        "pypi_315_simple_py3_none_any_deadbeef_local_override": {
+            "config_load": "@pypi//:config.bzl",
+            "dep_template": "@pypi//{name}:{target}",
+            "filename": "simple-0.0.3-py3-none-any.whl",
+            "index_url": "https://example.com",
+            "requirement": "simple==0.0.1",
+            "sha256": "",
+            "urls": ["file://custom_folder/simple-0.0.3-py3-none-any.whl"],
+        },
+    })
+
+_tests.append(_test_explicit_local_wheels)
+
 def hub_builder_test_suite(name):
     """Create the test suite.
 
diff --git a/tests/support/mocks/mocks.bzl b/tests/support/mocks/mocks.bzl
index 48f8f95830..b4398f3f45 100644
--- a/tests/support/mocks/mocks.bzl
+++ b/tests/support/mocks/mocks.bzl
@@ -11,12 +11,58 @@ def _path_new(path, mock_files = None):
         {type}`MockPath` A struct mocking a path object.
     """
     mock_files = mock_files or {}
-    return struct(
-        exists = path in mock_files,
-        basename = path.split("/")[-1],
-        dirname = "/".join(path.split("/")[:-1]),
-        _path = path,
-    )
+    path_str = str(path)
+    if path_str.startswith("@@//:"):
+        path_str = path_str[5:]
+    elif path_str.startswith("//:"):
+        path_str = path_str[3:]
+
+    parts = path_str.split("/") if path_str else []
+
+    current_struct = None
+
+    for i in range(len(parts) + 1):
+        sub_path = "/".join(parts[:i])
+        parent_struct = current_struct
+
+        def _get_child(child_name, p = sub_path):
+            child_path = "{}/{}".format(p, child_name) if p else str(child_name)
+            return _path_new(child_path, mock_files)
+
+        def _readdir(p = sub_path):
+            prefix = (p + "/") if p else ""
+            return [
+                _path_new(f, mock_files)
+                for f in mock_files
+                if f.startswith(prefix) and f != p
+            ]
+
+        def _exists(p = sub_path):
+            if p in mock_files:
+                return True
+            prefix = (p + "/") if p else ""
+            for f in mock_files:
+                if f.startswith(prefix):
+                    return True
+            return False
+
+        current_struct = struct(
+            exists = _exists(),
+            basename = parts[i - 1] if i > 0 else "",
+            dirname = parent_struct if i > 0 else struct(
+                exists = True,
+                basename = "",
+                dirname = "",
+                get_child = _get_child,
+                readdir = _readdir,
+                _path = "",
+            ),
+            get_child = _get_child,
+            readdir = _readdir,
+            _path = sub_path,
+        )
+
+    return current_struct
 
 def _file_new(short_path, *, path = None, is_source = True, owner = None):
     """Create a mock File object.
@@ -67,7 +113,7 @@ def _file_new(short_path, *, path = None, is_source = True, owner = None):
     return struct(
         path = path,
         basename = path.split("/")[-1],
-        dirname = "/".join(path.split("/")[:-1]),
+        dirname = _path_new("/".join(path.split("/")[:-1])),
         extension = path.split(".")[-1] if "." in path else "",
         is_source = is_source,
         owner = owner,
diff --git a/tests/support/mocks/mocks_tests.bzl b/tests/support/mocks/mocks_tests.bzl
index 1dc495b342..d5eb6c4a60 100644
--- a/tests/support/mocks/mocks_tests.bzl
+++ b/tests/support/mocks/mocks_tests.bzl
@@ -9,7 +9,7 @@ def _test_path(env):
     p1 = mocks.path("a/b/c", mock_files = {"a/b/c": "data"})
     env.expect.that_bool(p1.exists).equals(True)
     env.expect.that_str(p1.basename).equals("c")
-    env.expect.that_str(p1.dirname).equals("a/b")
+    env.expect.that_str(p1.dirname._path).equals("a/b")
     env.expect.that_str(p1._path).equals("a/b/c")
 
     p2 = mocks.path("d/e/f", mock_files = {})
@@ -23,7 +23,7 @@ def _test_file(env):
     env.expect.that_str(f1.path).equals("a/b.txt")
     env.expect.that_str(f1.short_path).equals("a/b.txt")
     env.expect.that_str(f1.basename).equals("b.txt")
-    env.expect.that_str(f1.dirname).equals("a")
+    env.expect.that_str(f1.dirname._path).equals("a")
     env.expect.that_str(f1.extension).equals("txt")
     env.expect.that_bool(f1.is_source).equals(True)
     env.expect.that_str(str(f1.owner)).equals(str(Label("//:mock")))

From 8310448b6acce1300def128dd9969b59ea227123 Mon Sep 17 00:00:00 2001
From: Yun Peng <pcloudy@google.com>
Date: Tue, 12 May 2026 13:51:38 +0000
Subject: [PATCH 02/11] Fix buildifier

---
 python/private/pypi/extension.bzl            |  6 ++--
 python/private/pypi/hub_builder.bzl          |  1 +
 tests/pypi/hub_builder/hub_builder_tests.bzl | 32 ++++++++++----------
 3 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
index 3280786933..314258b8bf 100644
--- a/python/private/pypi/extension.bzl
+++ b/python/private/pypi/extension.bzl
@@ -821,6 +821,9 @@ a string `"{os}_{arch}"` as the value here. You could also use `"{os}_{arch}_fre
 :::
 """,
         ),
+        "local_wheels": attr.string_dict(
+            doc = "Dictionary mapping package names to local wheel file paths relative to the workspace root.",
+        ),
         "whl_modifications": attr.label_keyed_string_dict(
             mandatory = False,
             doc = """\
@@ -828,9 +831,6 @@ A dict of labels to wheel names that is typically generated by the whl_modificat
 The labels are JSON config files describing the modifications.
 """,
         ),
-        "local_wheels": attr.string_dict(
-            doc = "Dictionary mapping package names to local wheel file paths relative to the workspace root.",
-        ),
     }, **ATTRS)
     attrs.update(AUTH_ATTRS)
 
diff --git a/python/private/pypi/hub_builder.bzl b/python/private/pypi/hub_builder.bzl
index c0a7233a49..4de9212a02 100644
--- a/python/private/pypi/hub_builder.bzl
+++ b/python/private/pypi/hub_builder.bzl
@@ -488,6 +488,7 @@ def _create_whl_repos(
         self: the builder.
         module_ctx: {type}`module_ctx`.
         pip_attr: {type}`struct` - the struct that comes from the tag class iteration.
+        is_root: {type}`bool` - whether the calling module is the root workspace.
         enable_pipstar_extract: {type}`bool` - enable the pipstar extraction or not.
     """
     logger = self._logger
diff --git a/tests/pypi/hub_builder/hub_builder_tests.bzl b/tests/pypi/hub_builder/hub_builder_tests.bzl
index 72b46d778b..28da68c672 100644
--- a/tests/pypi/hub_builder/hub_builder_tests.bzl
+++ b/tests/pypi/hub_builder/hub_builder_tests.bzl
@@ -853,9 +853,9 @@ def _test_local_wheel_override(env):
     builder.pip_parse(
         mocks.mctx(
             mock_files = {
-                "requirements.txt": "simple==0.0.1 --hash=sha256:deadbeef",
                 "MODULE.bazel": "",
                 "dist/simple-0.0.2-cp315-cp315-linux_x86_64.whl": "",
+                "requirements.txt": "simple==0.0.1 --hash=sha256:deadbeef",
             },
             os_name = "linux",
             arch_name = "x86_64",
@@ -901,9 +901,9 @@ def _test_local_wheel_override_ignored_if_not_root(env):
     builder.pip_parse(
         mocks.mctx(
             mock_files = {
-                "requirements.txt": "simple==0.0.1 --hash=sha256:deadbeef",
                 "MODULE.bazel": "",
                 "dist/simple-0.0.2-cp315-cp315-linux_x86_64.whl": "",
+                "requirements.txt": "simple==0.0.1 --hash=sha256:deadbeef",
             },
             os_name = "linux",
             arch_name = "x86_64",
@@ -1602,26 +1602,26 @@ _tests.append(_test_err_duplicate_repos)
 def _test_explicit_local_wheels(env):
     def mock_simpleapi_download(*_, **__):
         return {
-            "simple": struct(
+            "libtpu": struct(
                 whls = {
-                    "deadbeef": struct(
+                    "deadbaaf": struct(
                         yanked = None,
-                        filename = "simple-0.0.1-py3-none-any.whl",
-                        sha256 = "deadbeef",
-                        url = "example.com/simple-0.0.1.whl",
+                        filename = "libtpu-0.0.1-py3-none-any.whl",
+                        sha256 = "deadbaaf",
+                        url = "example.com/libtpu-0.0.1.whl",
                     ),
                 },
                 sdists = {},
                 sha256s_by_version = {},
                 index_url = "https://example.com",
             ),
-            "libtpu": struct(
+            "simple": struct(
                 whls = {
-                    "deadbaaf": struct(
+                    "deadbeef": struct(
                         yanked = None,
-                        filename = "libtpu-0.0.1-py3-none-any.whl",
-                        sha256 = "deadbaaf",
-                        url = "example.com/libtpu-0.0.1.whl",
+                        filename = "simple-0.0.1-py3-none-any.whl",
+                        sha256 = "deadbeef",
+                        url = "example.com/simple-0.0.1.whl",
                     ),
                 },
                 sdists = {},
@@ -1637,13 +1637,13 @@ def _test_explicit_local_wheels(env):
     builder.pip_parse(
         mocks.mctx(
             mock_files = {
+                "MODULE.bazel": "",
+                "custom_folder/libtpu-0.0.41.dev20260509+nightly-cp314-cp314t-manylinux_2_31_x86_64.whl": "",
+                "custom_folder/simple-0.0.3-py3-none-any.whl": "",
                 "requirements.txt": """\
 simple==0.0.1 --hash=sha256:deadbeef
 libtpu==0.0.1 --hash=sha256:deadbaaf
 """,
-                "MODULE.bazel": "",
-                "custom_folder/libtpu-0.0.41.dev20260509+nightly-cp314-cp314t-manylinux_2_31_x86_64.whl": "",
-                "custom_folder/simple-0.0.3-py3-none-any.whl": "",
             },
             os_name = "linux",
             arch_name = "x86_64",
@@ -1654,8 +1654,8 @@ libtpu==0.0.1 --hash=sha256:deadbaaf
             experimental_index_url = "https://example.com",
             requirements_lock = "requirements.txt",
             local_wheels = {
-                "simple": "custom_folder/simple-0.0.3-py3-none-any.whl",
                 "libtpu": "custom_folder/libtpu-*.whl",
+                "simple": "custom_folder/simple-0.0.3-py3-none-any.whl",
             },
         ),
         is_root = True,

From c34f7e8309addfed5b7c5f54c0108ca184be2dc0 Mon Sep 17 00:00:00 2001
From: Yun Peng <pcloudy@google.com>
Date: Tue, 12 May 2026 14:00:18 +0000
Subject: [PATCH 03/11] Fix buildifier

---
 python/private/pypi/extension.bzl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
index 314258b8bf..bb5721d8a5 100644
--- a/python/private/pypi/extension.bzl
+++ b/python/private/pypi/extension.bzl
@@ -761,6 +761,9 @@ hubs can be created, and each program can use its respective hub's targets.
 Targets from different hubs should not be used together.
 """,
         ),
+        "local_wheels": attr.string_dict(
+            doc = "Dictionary mapping package names to local wheel file paths relative to the workspace root.",
+        ),
         "parallel_download": attr.bool(
             doc = """\
 The flag allows to make use of parallel downloading feature in bazel 7.1 and above
@@ -821,9 +824,6 @@ a string `"{os}_{arch}"` as the value here. You could also use `"{os}_{arch}_fre
 :::
 """,
         ),
-        "local_wheels": attr.string_dict(
-            doc = "Dictionary mapping package names to local wheel file paths relative to the workspace root.",
-        ),
         "whl_modifications": attr.label_keyed_string_dict(
             mandatory = False,
             doc = """\

From f678e1120131303dff55fdd468d0fb9c8782691b Mon Sep 17 00:00:00 2001
From: Yun Peng <pcloudy@google.com>
Date: Tue, 12 May 2026 14:09:48 +0000
Subject: [PATCH 04/11] Fix workspace test

---
 tests/support/mocks/mocks.bzl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/support/mocks/mocks.bzl b/tests/support/mocks/mocks.bzl
index b4398f3f45..7d84bd24fc 100644
--- a/tests/support/mocks/mocks.bzl
+++ b/tests/support/mocks/mocks.bzl
@@ -14,6 +14,8 @@ def _path_new(path, mock_files = None):
     path_str = str(path)
     if path_str.startswith("@@//:"):
         path_str = path_str[5:]
+    elif path_str.startswith("@//:"):
+        path_str = path_str[4:]
     elif path_str.startswith("//:"):
         path_str = path_str[3:]
 

From 6506625b926efe320ff030ba8dbbf8d96c67a370 Mon Sep 17 00:00:00 2001
From: Yun Peng <pcloudy@google.com>
Date: Tue, 12 May 2026 14:17:49 +0000
Subject: [PATCH 05/11] Fix

---
 python/private/pypi/hub_builder.bzl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/private/pypi/hub_builder.bzl b/python/private/pypi/hub_builder.bzl
index 4de9212a02..b5e35727f0 100644
--- a/python/private/pypi/hub_builder.bzl
+++ b/python/private/pypi/hub_builder.bzl
@@ -333,7 +333,7 @@ def _add_whl_library(self, *, python_version, whl, repo):
                     if value
                 ])
             ))
-        return
+            return
     self._whl_libraries[repo_name] = repo.args
 
     mapping = self._whl_map.setdefault(whl.name, {})

From 1f8d43c8c0c3ecedb373b51a12b0253501d4c93c Mon Sep 17 00:00:00 2001
From: Yun Peng <pcloudy@google.com>
Date: Tue, 12 May 2026 14:43:34 +0000
Subject: [PATCH 06/11] Update doc

---
 python/private/pypi/extension.bzl | 27 ++++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
index bb5721d8a5..df54f73b85 100644
--- a/python/private/pypi/extension.bzl
+++ b/python/private/pypi/extension.bzl
@@ -581,8 +581,8 @@ Index metadata will be used to get `sha256` values for packages even if the
 Defaults to `https://pypi.org/simple`.
 
 :::{versionadded} 2.0.0
-This has been added as a replacement for 
-{obj}`pip.parse.experimental_index_url` and 
+This has been added as a replacement for
+{obj}`pip.parse.experimental_index_url` and
 {obj}`pip.parse.experimental_extra_index_urls`.
 :::
 """,
@@ -762,7 +762,28 @@ Targets from different hubs should not be used together.
 """,
         ),
         "local_wheels": attr.string_dict(
-            doc = "Dictionary mapping package names to local wheel file paths relative to the workspace root.",
+            doc = """\
+A dictionary mapping package names to local wheel file paths relative to the
+workspace root.
+
+Allows testing locally built wheels without modifying lockfiles or hosting a
+local index server.
+
+Keys are normalized package names (e.g. `my_package`). Values are paths relative
+to the workspace root.
+
+If the path contains a wildcard `*` (e.g. `dist/libtpu-*.whl`), matching `.whl`
+candidates are discovered and the newest version is selected by comparing
+basenames lexicographically.
+
+If a file is missing on disk or no files match the wildcard pattern, the
+override is silently ignored and Bazel falls back to the remote PyPI index.
+
+Overrides apply when bazel downloader is used and only take effect in the root module.
+
+:::{versionadded} 1.9.0
+:::
+""",
         ),
         "parallel_download": attr.bool(
             doc = """\

From 8f9d723967571a248ca1921aaf1f09b8c2baae52 Mon Sep 17 00:00:00 2001
From: Yun Peng <pcloudy@google.com>
Date: Wed, 13 May 2026 07:45:21 +0000
Subject: [PATCH 07/11] Update version

---
 python/private/pypi/extension.bzl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
index df54f73b85..3e295c87a7 100644
--- a/python/private/pypi/extension.bzl
+++ b/python/private/pypi/extension.bzl
@@ -781,7 +781,7 @@ override is silently ignored and Bazel falls back to the remote PyPI index.
 
 Overrides apply when bazel downloader is used and only take effect in the root module.
 
-:::{versionadded} 1.9.0
+:::{versionadded} 2.1.0
 :::
 """,
         ),

From 42296dc2ff77ba12b3bc0c97be9a3ee942c414c3 Mon Sep 17 00:00:00 2001
From: Yun Peng <pcloudy@google.com>
Date: Fri, 15 May 2026 14:44:00 +0000
Subject: [PATCH 08/11] Re-structure with reviewer feedback

---
 python/private/pypi/extension.bzl            |  34 ++--
 python/private/pypi/hub_builder.bzl          | 165 +++++++++----------
 python/private/pypi/parse_requirements.bzl   |  43 ++---
 tests/pypi/hub_builder/hub_builder_tests.bzl | 103 +-----------
 tests/support/mocks/mocks.bzl                |  16 ++
 5 files changed, 141 insertions(+), 220 deletions(-)

diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl
index 3e295c87a7..68e75fd441 100644
--- a/python/private/pypi/extension.bzl
+++ b/python/private/pypi/extension.bzl
@@ -761,29 +761,35 @@ hubs can be created, and each program can use its respective hub's targets.
 Targets from different hubs should not be used together.
 """,
         ),
-        "local_wheels": attr.string_dict(
+        "local_wheel_dir": attr.string(
             doc = """\
-A dictionary mapping package names to local wheel file paths relative to the
-workspace root.
+A path to a directory containing locally built wheels, relative to the workspace root.
 
-Allows testing locally built wheels without modifying lockfiles or hosting a
-local index server.
+Allows testing locally built wheels without modifying lockfiles or hosting a local index server.
 
-Keys are normalized package names (e.g. `my_package`). Values are paths relative
-to the workspace root.
+Must be used in conjunction with `local_wheel_pkgs` to specify which packages should be overridden.
+All child wheels in the directory belonging to the specified packages are scanned and selected to
+find the optimal wheel for the target platform.
 
-If the path contains a wildcard `*` (e.g. `dist/libtpu-*.whl`), matching `.whl`
-candidates are discovered and the newest version is selected by comparing
-basenames lexicographically.
-
-If a file is missing on disk or no files match the wildcard pattern, the
-override is silently ignored and Bazel falls back to the remote PyPI index.
+If the directory is missing on disk or no compatible wheel is found, the override is silently
+ignored and Bazel falls back to the remote PyPI index.
 
 Overrides apply when bazel downloader is used and only take effect in the root module.
 
-:::{versionadded} 2.1.0
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
+""",
+        ),
+        "local_wheel_pkgs": attr.string_list(
+            doc = """\
+A list of normalized package names (e.g. `my_package`) to override with wheels from `local_wheel_dir`.
+
+Must be used in conjunction with `local_wheel_dir`.
+
+:::{versionadded} VERSION_NEXT_FEATURE
 :::
 """,
+            default = [],
         ),
         "parallel_download": attr.bool(
             doc = """\
diff --git a/python/private/pypi/hub_builder.bzl b/python/private/pypi/hub_builder.bzl
index b5e35727f0..d0a70f2df7 100644
--- a/python/private/pypi/hub_builder.bzl
+++ b/python/private/pypi/hub_builder.bzl
@@ -11,8 +11,10 @@ load(":evaluate_markers.bzl", "evaluate_markers")
 load(":parse_requirements.bzl", "parse_requirements")
 load(":pep508_env.bzl", "env")
 load(":pep508_evaluate.bzl", "evaluate")
+load(":parse_whl_name.bzl", "parse_whl_name")
 load(":python_tag.bzl", "python_tag")
 load(":requirements_files_by_platform.bzl", "requirements_files_by_platform")
+load(":version_from_filename.bzl", "version_from_filename")
 load(":whl_config_setting.bzl", "whl_config_setting")
 load(":whl_repo_name.bzl", "pypi_repo_name", "whl_repo_name")
 
@@ -181,7 +183,7 @@ def _pip_parse(self, module_ctx, pip_attr, is_root = False):
         ))
         return
 
-    _set_get_index_urls(self, pip_attr)
+    _set_get_index_urls(self, module_ctx, pip_attr, is_root = is_root)
     self._platforms[python_version] = _platforms(
         module_ctx,
         python_version = full_python_version,
@@ -350,7 +352,7 @@ def _add_whl_library(self, *, python_version, whl, repo):
 
 ### end of setters, below we have various functions to implement the public methods
 
-def _set_get_index_urls(self, pip_attr):
+def _set_get_index_urls(self, module_ctx, pip_attr, is_root = False):
     default_index_url = pip_attr.experimental_index_url or self._config.index_url
     default_extra_index_urls = pip_attr.experimental_extra_index_urls or []
 
@@ -364,28 +366,76 @@ def _set_get_index_urls(self, pip_attr):
         normalize_name(s): False
         for s in pip_attr.simpleapi_skip
     })
-    self._get_index_urls[python_version] = lambda ctx, distributions, *, index_url = None, extra_index_urls = None: self._simpleapi_download_fn(
-        ctx,
-        attr = struct(
-            index_url = (index_url or default_index_url).rstrip("/"),
-            extra_index_urls = [
-                x.rstrip("/")
-                for x in (extra_index_urls or default_extra_index_urls)
-            ],
-            index_url_overrides = pip_attr.experimental_index_url_overrides or {},
-            sources = {
-                d: versions
-                for d, versions in distributions.items()
-                if _use_downloader(self, python_version, d)
-            },
-            envsubst = pip_attr.envsubst,
-            # Auth related info
-            netrc = self._config.netrc or pip_attr.netrc,
-            auth_patterns = self._config.auth_patterns or pip_attr.auth_patterns,
-        ),
-        cache = self._simpleapi_cache,
-        parallel_download = pip_attr.parallel_download,
-    )
+
+    def _download_wrapper(ctx, distributions, *, index_url = None, extra_index_urls = None):
+        res = self._simpleapi_download_fn(
+            ctx,
+            attr = struct(
+                index_url = (index_url or default_index_url).rstrip("/"),
+                extra_index_urls = [
+                    x.rstrip("/")
+                    for x in (extra_index_urls or default_extra_index_urls)
+                ],
+                index_url_overrides = pip_attr.experimental_index_url_overrides or {},
+                sources = {
+                    d: versions
+                    for d, versions in distributions.items()
+                    if _use_downloader(self, python_version, d)
+                },
+                envsubst = pip_attr.envsubst,
+                # Auth related info
+                netrc = self._config.netrc or pip_attr.netrc,
+                auth_patterns = self._config.auth_patterns or pip_attr.auth_patterns,
+            ),
+            cache = self._simpleapi_cache,
+            parallel_download = pip_attr.parallel_download,
+        )
+
+        if not is_root or not getattr(pip_attr, "local_wheel_dir", None) or not getattr(pip_attr, "local_wheel_pkgs", None):
+            return res
+
+        workspace_root = module_ctx.path(Label("@@//:MODULE.bazel")).dirname
+        target_dir = workspace_root.get_child(pip_attr.local_wheel_dir)
+        if not target_dir.exists or not getattr(target_dir, "is_dir", False):
+            return res
+
+        candidates = target_dir.readdir()
+        override_pkgs = {normalize_name(p): True for p in pip_attr.local_wheel_pkgs}
+
+        wheels_by_pkg = {}
+        for candidate in candidates:
+            if not candidate.basename.endswith(".whl"):
+                continue
+            parsed = parse_whl_name(candidate.basename)
+            norm_name = normalize_name(parsed.distribution)
+            if norm_name in override_pkgs and norm_name in distributions:
+                wheels_by_pkg.setdefault(norm_name, []).append(candidate)
+
+        for norm_name, matched_wheels in wheels_by_pkg.items():
+            local_dists = []
+            for wheel_path in matched_wheels:
+                dist = struct(
+                    filename = wheel_path.basename,
+                    version = version_from_filename(wheel_path.basename),
+                    url = "file://" + (wheel_path._path if hasattr(wheel_path, "_path") else str(wheel_path)),
+                    sha256 = "",
+                    metadata_sha256 = "",
+                    metadata_url = "",
+                    yanked = None,
+                )
+                local_dists.append(dist)
+
+            res[norm_name] = struct(
+                sdists = {},
+                whls = {d.filename: d for d in local_dists},
+                sha256s_by_version = {},
+                index_url = "file://" + (workspace_root._path if hasattr(workspace_root, "_path") else str(workspace_root)),
+                local_override_whls = local_dists,
+            )
+
+        return res
+
+    self._get_index_urls[python_version] = _download_wrapper
     return True
 
 def _detect_interpreter(self, pip_attr):
@@ -535,10 +585,7 @@ def _create_whl_repos(
 
     interpreter = _detect_interpreter(self, pip_attr)
 
-    local_wheels = _collect_local_wheels(module_ctx, pip_attr, is_root = is_root)
-
     for whl in requirements_by_platform:
-        local_wheel = local_wheels.get(whl.name)
         whl_library_args = common_args | _whl_library_args(
             self,
             whl = whl,
@@ -556,7 +603,6 @@ def _create_whl_repos(
                 python_version = _major_minor_version(pip_attr.python_version),
                 is_multiple_versions = whl.is_multiple_versions,
                 interpreter = interpreter,
-                local_wheel = local_wheel,
                 enable_pipstar_extract = enable_pipstar_extract,
             )
             _add_whl_library(
@@ -632,7 +678,6 @@ def _whl_repo(
         python_version,
         use_downloader,
         interpreter,
-        local_wheel = None,
         enable_pipstar_extract = False):
     args = dict(whl_library_args)
     args["requirement"] = src.requirement_line
@@ -691,12 +736,8 @@ def _whl_repo(
     target_platforms = src.target_platforms if is_multiple_versions else []
     repo_name = whl_repo_name(src.filename, src.sha256, *target_platforms)
 
-    if local_wheel:
+    if src.url.startswith("file://"):
         repo_name += "_local_override"
-        path_str = local_wheel._path if hasattr(local_wheel, "_path") else str(local_wheel)
-        args["urls"] = ["file://" + path_str]
-        args["filename"] = local_wheel.basename
-        args["sha256"] = ""
 
     return struct(
         repo_name = repo_name,
@@ -713,60 +754,4 @@ def _use_downloader(self, python_version, whl_name):
         self._get_index_urls.get(python_version) != None,
     )
 
-def _collect_local_wheels(module_ctx, pip_attr, is_root = False):
-    if not is_root:
-        return {}
-
-    wheels = {}
-    explicit_wheels = getattr(pip_attr, "local_wheels", None)
-    if not explicit_wheels:
-        return wheels
-
-    workspace_root = module_ctx.path(Label("@@//:MODULE.bazel")).dirname
 
-    for pkg_name, wheel_path_str in explicit_wheels.items():
-        norm_name = normalize_name(pkg_name)
-        if "*" not in wheel_path_str:
-            wheel_path = workspace_root.get_child(wheel_path_str)
-            if wheel_path.exists:
-                wheels[norm_name] = wheel_path
-        else:
-            last_slash = wheel_path_str.rfind("/")
-            if last_slash >= 0:
-                dir_part = wheel_path_str[:last_slash]
-                pattern = wheel_path_str[last_slash + 1:]
-            else:
-                dir_part = ""
-                pattern = wheel_path_str
-
-            matched_wheel = None
-            target_dir = workspace_root.get_child(dir_part) if dir_part else workspace_root
-            if target_dir.exists:
-                candidates = target_dir.readdir()
-            else:
-                candidates = []
-
-            for candidate in candidates:
-                if not candidate.basename.endswith(".whl"):
-                    continue
-                if _wildcard_match(candidate.basename, pattern):
-                    if not matched_wheel or matched_wheel.basename < candidate.basename:
-                        matched_wheel = candidate
-
-            if matched_wheel:
-                wheels[norm_name] = matched_wheel
-
-    return wheels
-
-def _wildcard_match(name, pattern):
-    if pattern.startswith("*") and pattern.endswith("*"):
-        return name.find(pattern[1:-1]) >= 0
-    elif pattern.startswith("*"):
-        return name.endswith(pattern[1:])
-    elif pattern.endswith("*"):
-        return name.startswith(pattern[:-1])
-    elif "*" in pattern:
-        parts = pattern.split("*", 1)
-        return name.startswith(parts[0]) and name.endswith(parts[1]) and len(name) >= len(parts[0]) + len(parts[1])
-    else:
-        return name == pattern
diff --git a/python/private/pypi/parse_requirements.bzl b/python/private/pypi/parse_requirements.bzl
index 07d0c0989e..a1347898b4 100644
--- a/python/private/pypi/parse_requirements.bzl
+++ b/python/private/pypi/parse_requirements.bzl
@@ -429,29 +429,32 @@ def _add_dists(*, requirement, index_urls, target_platform, logger = None):
     whls = []
     sdist = None
 
-    # First try to find distributions by SHA256 if provided
-    shas_to_use = requirement.srcs.shas
-    if not shas_to_use:
-        version = requirement.srcs.version
-        shas_to_use = index_urls.sha256s_by_version.get(version, [])
-        logger.warn(lambda: "requirement file has been generated without hashes, will use all hashes for the given version {} that could find on the index:\n    {}".format(version, shas_to_use))
-
-    for sha256 in shas_to_use:
-        # For now if the artifact is marked as yanked we just ignore it.
-        #
-        # See https://packaging.python.org/en/latest/specifications/simple-repository-api/#adding-yank-support-to-the-simple-api
+    if hasattr(index_urls, "local_override_whls"):
+        whls = index_urls.local_override_whls
+    else:
+        # First try to find distributions by SHA256 if provided
+        shas_to_use = requirement.srcs.shas
+        if not shas_to_use:
+            version = requirement.srcs.version
+            shas_to_use = index_urls.sha256s_by_version.get(version, [])
+            logger.warn(lambda: "requirement file has been generated without hashes, will use all hashes for the given version {} that could find on the index:\n    {}".format(version, shas_to_use))
+
+        for sha256 in shas_to_use:
+            # For now if the artifact is marked as yanked we just ignore it.
+            #
+            # See https://packaging.python.org/en/latest/specifications/simple-repository-api/#adding-yank-support-to-the-simple-api
 
-        maybe_whl = index_urls.whls.get(sha256)
-        if maybe_whl and maybe_whl.yanked == None:
-            whls.append(maybe_whl)
-            continue
+            maybe_whl = index_urls.whls.get(sha256)
+            if maybe_whl and maybe_whl.yanked == None:
+                whls.append(maybe_whl)
+                continue
 
-        maybe_sdist = index_urls.sdists.get(sha256)
-        if maybe_sdist and maybe_sdist.yanked == None:
-            sdist = maybe_sdist
-            continue
+            maybe_sdist = index_urls.sdists.get(sha256)
+            if maybe_sdist and maybe_sdist.yanked == None:
+                sdist = maybe_sdist
+                continue
 
-        logger.warn(lambda: "Could not find a whl or an sdist with sha256={}".format(sha256))
+            logger.warn(lambda: "Could not find a whl or an sdist with sha256={}".format(sha256))
 
     yanked = {}
     for dist in whls + [sdist]:
diff --git a/tests/pypi/hub_builder/hub_builder_tests.bzl b/tests/pypi/hub_builder/hub_builder_tests.bzl
index 28da68c672..0e8b6b21fe 100644
--- a/tests/pypi/hub_builder/hub_builder_tests.bzl
+++ b/tests/pypi/hub_builder/hub_builder_tests.bzl
@@ -865,9 +865,8 @@ def _test_local_wheel_override(env):
             python_version = "3.15",
             experimental_index_url = "https://example.com",
             requirements_lock = "requirements.txt",
-            local_wheels = {
-                "simple": "dist/simple-0.0.2-cp315-cp315-linux_x86_64.whl",
-            },
+            local_wheel_dir = "dist",
+            local_wheel_pkgs = ["simple"],
         ),
         is_root = True,
     )
@@ -876,17 +875,17 @@ def _test_local_wheel_override(env):
     pypi.exposed_packages().contains_exactly(["simple"])
     pypi.whl_map().contains_exactly({
         "simple": {
-            "pypi_315_simple_py3_none_any_deadbeef_local_override": [
+            "pypi_315_simple_0_0_2_cp315_cp315_linux_x86_64_local_override": [
                 whl_config_setting(version = "3.15", target_platforms = ["cp315_linux_x86_64"]),
             ],
         },
     })
     pypi.whl_libraries().contains_exactly({
-        "pypi_315_simple_py3_none_any_deadbeef_local_override": {
+        "pypi_315_simple_0_0_2_cp315_cp315_linux_x86_64_local_override": {
             "config_load": "@pypi//:config.bzl",
             "dep_template": "@pypi//{name}:{target}",
             "filename": "simple-0.0.2-cp315-cp315-linux_x86_64.whl",
-            "index_url": "https://example.com",
+            "index_url": "file://",
             "requirement": "simple==0.0.1",
             "sha256": "",
             "urls": ["file://dist/simple-0.0.2-cp315-cp315-linux_x86_64.whl"],
@@ -912,9 +911,8 @@ def _test_local_wheel_override_ignored_if_not_root(env):
             hub_name = "pypi",
             python_version = "3.15",
             requirements_lock = "requirements.txt",
-            local_wheels = {
-                "simple": "dist/simple-0.0.2-cp315-cp315-linux_x86_64.whl",
-            },
+            local_wheel_dir = "dist",
+            local_wheel_pkgs = ["simple"],
         ),
         is_root = False,
     )
@@ -1599,93 +1597,6 @@ Attempting to create a duplicate library pypi_315_foo for foo with different arg
 
 _tests.append(_test_err_duplicate_repos)
 
-def _test_explicit_local_wheels(env):
-    def mock_simpleapi_download(*_, **__):
-        return {
-            "libtpu": struct(
-                whls = {
-                    "deadbaaf": struct(
-                        yanked = None,
-                        filename = "libtpu-0.0.1-py3-none-any.whl",
-                        sha256 = "deadbaaf",
-                        url = "example.com/libtpu-0.0.1.whl",
-                    ),
-                },
-                sdists = {},
-                sha256s_by_version = {},
-                index_url = "https://example.com",
-            ),
-            "simple": struct(
-                whls = {
-                    "deadbeef": struct(
-                        yanked = None,
-                        filename = "simple-0.0.1-py3-none-any.whl",
-                        sha256 = "deadbeef",
-                        url = "example.com/simple-0.0.1.whl",
-                    ),
-                },
-                sdists = {},
-                sha256s_by_version = {},
-                index_url = "https://example.com",
-            ),
-        }
-
-    builder = hub_builder(
-        env,
-        simpleapi_download_fn = mock_simpleapi_download,
-    )
-    builder.pip_parse(
-        mocks.mctx(
-            mock_files = {
-                "MODULE.bazel": "",
-                "custom_folder/libtpu-0.0.41.dev20260509+nightly-cp314-cp314t-manylinux_2_31_x86_64.whl": "",
-                "custom_folder/simple-0.0.3-py3-none-any.whl": "",
-                "requirements.txt": """\
-simple==0.0.1 --hash=sha256:deadbeef
-libtpu==0.0.1 --hash=sha256:deadbaaf
-""",
-            },
-            os_name = "linux",
-            arch_name = "x86_64",
-        ),
-        _parse(
-            hub_name = "pypi",
-            python_version = "3.15",
-            experimental_index_url = "https://example.com",
-            requirements_lock = "requirements.txt",
-            local_wheels = {
-                "libtpu": "custom_folder/libtpu-*.whl",
-                "simple": "custom_folder/simple-0.0.3-py3-none-any.whl",
-            },
-        ),
-        is_root = True,
-    )
-    pypi = builder.build()
-
-    pypi.exposed_packages().contains_exactly(["libtpu", "simple"])
-    pypi.whl_libraries().contains_exactly({
-        "pypi_315_libtpu_py3_none_any_deadbaaf_local_override": {
-            "config_load": "@pypi//:config.bzl",
-            "dep_template": "@pypi//{name}:{target}",
-            "filename": "libtpu-0.0.41.dev20260509+nightly-cp314-cp314t-manylinux_2_31_x86_64.whl",
-            "index_url": "https://example.com",
-            "requirement": "libtpu==0.0.1",
-            "sha256": "",
-            "urls": ["file://custom_folder/libtpu-0.0.41.dev20260509+nightly-cp314-cp314t-manylinux_2_31_x86_64.whl"],
-        },
-        "pypi_315_simple_py3_none_any_deadbeef_local_override": {
-            "config_load": "@pypi//:config.bzl",
-            "dep_template": "@pypi//{name}:{target}",
-            "filename": "simple-0.0.3-py3-none-any.whl",
-            "index_url": "https://example.com",
-            "requirement": "simple==0.0.1",
-            "sha256": "",
-            "urls": ["file://custom_folder/simple-0.0.3-py3-none-any.whl"],
-        },
-    })
-
-_tests.append(_test_explicit_local_wheels)
-
 def hub_builder_test_suite(name):
     """Create the test suite.
 
diff --git a/tests/support/mocks/mocks.bzl b/tests/support/mocks/mocks.bzl
index 7d84bd24fc..f3632ec808 100644
--- a/tests/support/mocks/mocks.bzl
+++ b/tests/support/mocks/mocks.bzl
@@ -3,6 +3,11 @@
 def _path_new(path, mock_files = None):
     """Create a mock path object.
 
+    Simulates a hierarchical directory structure by dynamically building recursive
+    struct representations for each parent directory segment. This allows calling
+    `.dirname`, `.get_child()`, `.readdir()`, and `.exists` on mock path instances,
+    which is essential for testing path traversal in starlark tests.
+
     Args:
         path: {type}`string` The path string.
         mock_files: {type}`dict[string, string]` A dict of mocked files.
@@ -48,11 +53,22 @@ def _path_new(path, mock_files = None):
                     return True
             return False
 
+        def _is_dir(p = sub_path):
+            if p == "":
+                return True
+            prefix = (p + "/") if p else ""
+            for f in mock_files:
+                if f.startswith(prefix) and f != p:
+                    return True
+            return False
+
         current_struct = struct(
             exists = _exists(),
+            is_dir = _is_dir(),
             basename = parts[i - 1] if i > 0 else "",
             dirname = parent_struct if i > 0 else struct(
                 exists = True,
+                is_dir = True,
                 basename = "",
                 dirname = "",
                 get_child = _get_child,

From 61a995b9d98b96f236250ce23945570fd7b45184 Mon Sep 17 00:00:00 2001
From: Yun Peng <pcloudy@google.com>
Date: Fri, 15 May 2026 15:02:49 +0000
Subject: [PATCH 09/11] Refactor

---
 python/private/pypi/hub_builder.bzl | 97 +++++++++++++++++------------
 1 file changed, 56 insertions(+), 41 deletions(-)

diff --git a/python/private/pypi/hub_builder.bzl b/python/private/pypi/hub_builder.bzl
index d0a70f2df7..dda33ada18 100644
--- a/python/private/pypi/hub_builder.bzl
+++ b/python/private/pypi/hub_builder.bzl
@@ -391,49 +391,10 @@ def _set_get_index_urls(self, module_ctx, pip_attr, is_root = False):
             parallel_download = pip_attr.parallel_download,
         )
 
-        if not is_root or not getattr(pip_attr, "local_wheel_dir", None) or not getattr(pip_attr, "local_wheel_pkgs", None):
+        if not is_root:
             return res
 
-        workspace_root = module_ctx.path(Label("@@//:MODULE.bazel")).dirname
-        target_dir = workspace_root.get_child(pip_attr.local_wheel_dir)
-        if not target_dir.exists or not getattr(target_dir, "is_dir", False):
-            return res
-
-        candidates = target_dir.readdir()
-        override_pkgs = {normalize_name(p): True for p in pip_attr.local_wheel_pkgs}
-
-        wheels_by_pkg = {}
-        for candidate in candidates:
-            if not candidate.basename.endswith(".whl"):
-                continue
-            parsed = parse_whl_name(candidate.basename)
-            norm_name = normalize_name(parsed.distribution)
-            if norm_name in override_pkgs and norm_name in distributions:
-                wheels_by_pkg.setdefault(norm_name, []).append(candidate)
-
-        for norm_name, matched_wheels in wheels_by_pkg.items():
-            local_dists = []
-            for wheel_path in matched_wheels:
-                dist = struct(
-                    filename = wheel_path.basename,
-                    version = version_from_filename(wheel_path.basename),
-                    url = "file://" + (wheel_path._path if hasattr(wheel_path, "_path") else str(wheel_path)),
-                    sha256 = "",
-                    metadata_sha256 = "",
-                    metadata_url = "",
-                    yanked = None,
-                )
-                local_dists.append(dist)
-
-            res[norm_name] = struct(
-                sdists = {},
-                whls = {d.filename: d for d in local_dists},
-                sha256s_by_version = {},
-                index_url = "file://" + (workspace_root._path if hasattr(workspace_root, "_path") else str(workspace_root)),
-                local_override_whls = local_dists,
-            )
-
-        return res
+        return _inject_local_wheels(module_ctx, pip_attr, distributions, res)
 
     self._get_index_urls[python_version] = _download_wrapper
     return True
@@ -754,4 +715,58 @@ def _use_downloader(self, python_version, whl_name):
         self._get_index_urls.get(python_version) != None,
     )
 
+def _inject_local_wheels(module_ctx, pip_attr, distributions, res):
+    """Inject local wheel overrides into the SimpleAPI download results.
+
+    Args:
+        module_ctx: {type}`module_ctx` The module context.
+        pip_attr: {type}`struct` The pip.parse attribute struct.
+        distributions: {type}`dict` The requested distributions map.
+        res: {type}`dict` The SimpleAPI download results dict.
+
+    Returns:
+        {type}`dict` The modified SimpleAPI download results dict.
+    """
+    if not getattr(pip_attr, "local_wheel_dir", None) or not getattr(pip_attr, "local_wheel_pkgs", None):
+        return res
+
+    workspace_root = module_ctx.path(Label("@@//:MODULE.bazel")).dirname
+    target_dir = workspace_root.get_child(pip_attr.local_wheel_dir)
+    if not target_dir.exists or not getattr(target_dir, "is_dir", False):
+        return res
+
+    candidates = target_dir.readdir()
+    override_pkgs = {normalize_name(p): True for p in pip_attr.local_wheel_pkgs}
+
+    wheels_by_pkg = {}
+    for candidate in candidates:
+        if not candidate.basename.endswith(".whl"):
+            continue
+        parsed = parse_whl_name(candidate.basename)
+        norm_name = normalize_name(parsed.distribution)
+        if norm_name in override_pkgs and norm_name in distributions:
+            wheels_by_pkg.setdefault(norm_name, []).append(candidate)
+
+    for norm_name, matched_wheels in wheels_by_pkg.items():
+        local_dists = []
+        for wheel_path in matched_wheels:
+            dist = struct(
+                filename = wheel_path.basename,
+                version = version_from_filename(wheel_path.basename),
+                url = "file://" + (wheel_path._path if hasattr(wheel_path, "_path") else str(wheel_path)),
+                sha256 = "",
+                metadata_sha256 = "",
+                metadata_url = "",
+                yanked = None,
+            )
+            local_dists.append(dist)
+
+        res[norm_name] = struct(
+            sdists = {},
+            whls = {d.filename: d for d in local_dists},
+            sha256s_by_version = {},
+            index_url = "file://" + (workspace_root._path if hasattr(workspace_root, "_path") else str(workspace_root)),
+            local_override_whls = local_dists,
+        )
 
+    return res

From 8aeb3b33b11bc8c42161c8174cc1fc5c614fc6da Mon Sep 17 00:00:00 2001
From: Yun Peng <pcloudy@google.com>
Date: Fri, 15 May 2026 15:07:52 +0000
Subject: [PATCH 10/11] Remove is_root

---
 python/private/pypi/hub_builder.bzl | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/python/private/pypi/hub_builder.bzl b/python/private/pypi/hub_builder.bzl
index dda33ada18..131a39c336 100644
--- a/python/private/pypi/hub_builder.bzl
+++ b/python/private/pypi/hub_builder.bzl
@@ -196,7 +196,6 @@ def _pip_parse(self, module_ctx, pip_attr, is_root = False):
         self,
         module_ctx,
         pip_attr = pip_attr,
-        is_root = is_root,
         enable_pipstar_extract = bool(self._config.enable_pipstar_extract or self._get_index_urls.get(pip_attr.python_version)),
     )
 
@@ -491,7 +490,6 @@ def _create_whl_repos(
         module_ctx,
         *,
         pip_attr,
-        is_root = False,
         enable_pipstar_extract = False):
     """create all of the whl repositories
 
@@ -499,7 +497,6 @@ def _create_whl_repos(
         self: the builder.
         module_ctx: {type}`module_ctx`.
         pip_attr: {type}`struct` - the struct that comes from the tag class iteration.
-        is_root: {type}`bool` - whether the calling module is the root workspace.
         enable_pipstar_extract: {type}`bool` - enable the pipstar extraction or not.
     """
     logger = self._logger

From 0cc002170ff3cee4d2413643a7bd955b08fddbb9 Mon Sep 17 00:00:00 2001
From: Yun Peng <pcloudy@google.com>
Date: Fri, 15 May 2026 15:15:15 +0000
Subject: [PATCH 11/11] Fix buildifier error

---
 python/private/pypi/hub_builder.bzl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/private/pypi/hub_builder.bzl b/python/private/pypi/hub_builder.bzl
index 131a39c336..5f5e014b1a 100644
--- a/python/private/pypi/hub_builder.bzl
+++ b/python/private/pypi/hub_builder.bzl
@@ -9,9 +9,9 @@ load("//python/private:version_label.bzl", "version_label")
 load(":attrs.bzl", "use_isolated")
 load(":evaluate_markers.bzl", "evaluate_markers")
 load(":parse_requirements.bzl", "parse_requirements")
+load(":parse_whl_name.bzl", "parse_whl_name")
 load(":pep508_env.bzl", "env")
 load(":pep508_evaluate.bzl", "evaluate")
-load(":parse_whl_name.bzl", "parse_whl_name")
 load(":python_tag.bzl", "python_tag")
 load(":requirements_files_by_platform.bzl", "requirements_files_by_platform")
 load(":version_from_filename.bzl", "version_from_filename")