feat(pypi): add native local wheel override support to bzlmod pip.parse

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.
@@ -580,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`.
 :::
 """,
@@ -760,6 +761,36 @@ hubs can be created, and each program can use its respective hub's targets.
 Targets from different hubs should not be used together.
 """,
         ),
+        "local_wheel_dir": attr.string(
+            doc = """\
+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.
+
+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 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} 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 = """\
 The flag allows to make use of parallel downloading feature in bazel 7.1 and above

python/private/pypi/hub_builder.bzl

@@ -9,10 +9,12 @@ 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(":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")
 
@@ -147,7 +149,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((
@@ -181,7 +183,7 @@ def _pip_parse(self, module_ctx, pip_attr):
         ))
         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,
@@ -349,7 +351,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 []
 
@@ -363,28 +365,37 @@ 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:
+            return res
+
+        return _inject_local_wheels(module_ctx, pip_attr, distributions, res)
+
+    self._get_index_urls[python_version] = _download_wrapper
     return True
 
 def _detect_interpreter(self, pip_attr):
@@ -681,9 +692,13 @@ 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 src.url.startswith("file://"):
+        repo_name += "_local_override"
 
     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,59 @@ def _use_downloader(self, python_version, whl_name):
         normalize_name(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

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]:

tests/pypi/hub_builder/hub_builder_tests.bzl

@@ -828,6 +828,108 @@ 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 = {
+                "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",
+        ),
+        _parse(
+            hub_name = "pypi",
+            python_version = "3.15",
+            experimental_index_url = "https://example.com",
+            requirements_lock = "requirements.txt",
+            local_wheel_dir = "dist",
+            local_wheel_pkgs = ["simple"],
+        ),
+        is_root = True,
+    )
+    pypi = builder.build()
+
+    pypi.exposed_packages().contains_exactly(["simple"])
+    pypi.whl_map().contains_exactly({
+        "simple": {
+            "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_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": "file://",
+            "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 = {
+                "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",
+        ),
+        _parse(
+            hub_name = "pypi",
+            python_version = "3.15",
+            requirements_lock = "requirements.txt",
+            local_wheel_dir = "dist",
+            local_wheel_pkgs = ["simple"],
+        ),
+        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(