diff --git a/python/private/pypi/extension.bzl b/python/private/pypi/extension.bzl index 2a8ace28f2..3e295c87a7 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. @@ -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`. ::: """, @@ -758,6 +759,30 @@ is not required. Each hub is a separate resolution of pip dependencies. This means if different programs need different versions of some library, separate 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 = """\ +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} 2.1.0 +::: """, ), "parallel_download": attr.bool( diff --git a/python/private/pypi/hub_builder.bzl b/python/private/pypi/hub_builder.bzl index 1bce648dce..b5e35727f0 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)), ) @@ -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 @@ -486,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 @@ -532,7 +535,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 +556,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 +632,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 +689,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 +712,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..28da68c672 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 = { + "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_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 = { + "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_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 { + "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 48f8f95830..7d84bd24fc 100644 --- a/tests/support/mocks/mocks.bzl +++ b/tests/support/mocks/mocks.bzl @@ -11,12 +11,60 @@ 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[4:] + 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 +115,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")))