Add an override module repos feature

.github/workflows/build.yaml

@@ -135,6 +135,10 @@ jobs:
       run: |
         bazel test //... --noenable_bzlmod
       working-directory: private/examples/multiplatform_py_test
+    - name: Test module repo overrides
+      run: |
+        bazel test //...
+      working-directory: private/tests/override_module_repos/overrider
 
   bazel_lint:
     runs-on: ubuntu-latest

.gitignore

@@ -13,3 +13,7 @@ dist
 user.bazelrc
 /.github/ci.bazelrc
 /private/examples/multiplatform_py_test/bazel-*
+
+# Ignore the module lock for now, since the sdist repo rules incorrectly record the calling
+# platform's Python intepreter as the _only_ interprete in the lock.
+MODULE.bazel.lock

MODULE.bazel

@@ -196,7 +196,7 @@ dev_requirements.package_annotation(
 
 # Use a separate extension since the @req_compile_find_links_test is generated by an extension
 # that also depend upon //extensions:python.bzl
-find_links_requirements = use_extension("//private/tests/find_links:extension.bzl", "requirements")
+find_links_requirements = use_extension("//private/tests/find_links:extension.bzl", "requirements", dev_dependency = True)
 find_links_requirements.parse(
     name = "req_compile_test_find_links",
     requirements_lock = "@req_compile_find_links_test//:requirements.txt",

extensions/python.bzl

@@ -15,6 +15,8 @@ def _requirements_impl(ctx):
     """Process annotations and parse tags."""
     annotations = {}
 
+    override_module_repos = {}
+
     # Gather all annotations first.
     for mod in ctx.modules:
         for annotation in mod.tags.package_annotation:
@@ -31,6 +33,14 @@ def _requirements_impl(ctx):
                 deps_excludes = annotation.deps_excludes,
                 patches = annotation.patches,
             )
+        for parse in mod.tags.parse:
+            if mod.is_root and parse.override_module_repos:
+                for override_mod, override_repos in parse.override_module_repos.items():
+                    if override_mod in override_module_repos:
+                        fail("Module {} already has overrides from this module.".format(
+                            override_mod,
+                        ))
+                override_module_repos[override_mod] = struct(hub_name = parse.name, override_repos = list(override_repos))
 
     # Create hubs for each parse tag.
     for mod in ctx.modules:
@@ -48,6 +58,17 @@ def _requirements_impl(ctx):
             elif ctx.os.name.startswith("windows"):
                 interpreter = parse.interpreter_windows
 
+            if mod.name in override_module_repos and parse.name in override_module_repos[mod.name].override_repos:
+                py_requirements_repository(
+                    name = parse.name,
+                    hub_name = override_module_repos[mod.name].hub_name,
+                    requirements_lock = parse.requirements_lock,
+                    requirements_locks = parse.requirements_locks,
+                    interpreter = interpreter,
+                )
+                override_module_repos[mod.name].override_repos.remove(parse.name)
+                continue
+
             py_requirements_repository(
                 name = parse.name,
                 hub_name = parse.name,
@@ -67,6 +88,11 @@ def _requirements_impl(ctx):
                     spoke_prefix += "_" + defs_id
                 create_spoke_repos(spoke_prefix, data.packages, interpreter)
 
+    for mod in sorted(override_module_repos):
+        repos = override_module_repos[mod].override_repos
+        if repos:
+            fail("Module \"{}\" does not create repos \"{}\"".format(mod, "\", \"".join(sorted(repos))))
+
 _annotation = tag_class(
     doc = """\
 A tag representing a annotation editing a Python package.
@@ -140,6 +166,20 @@ hub repository named "pip_deps".
         "requirements_locks": attr.label_keyed_string_dict(
             doc = "A dictionary mapping platform to requirement lock files.",
         ),
+        "override_module_repos": attr.string_list_dict(
+            doc = """\
+Mapping of module name to list of repos that should be overridden by this hub. The repos
+must be those that are expected to be created by this module extension.
+
+The overridden hub will attempt to map all of its requirements to the root module's hub, meaning
+the root module hub must be a superset of the overridden hub.
+
+This attribute is intended to have the root module coordinate all Python packages such
+that Python libraries from dependencies can be safely imported into the same interpreter.
+Do not override repos for libraries that will never be mixed. To inject Python dependencies
+for use in most child modules, a custom toolchain type is most appropriate.""",
+            default = {},
+        ),
     },
 )
 

private/examples/multiplatform_py_test/MODULE.bazel

@@ -1,3 +1,8 @@
+module(
+    name = "multiplatform_py_test",
+    version = "1.0",
+)
+
 bazel_dep(name = "rules_python", version = "0.34.0")
 bazel_dep(name = "platforms", version = "0.0.9")
 bazel_dep(name = "rules_req_compile", version = "0.0.0")

private/tests/override_module_repos/overridee/BUILD.bazel

@@ -0,0 +1,7 @@
+load("@rules_python//python:defs.bzl", "py_library")
+
+py_library(
+    name = "test",
+    visibility = ["//visibility:public"],
+    deps = ["@pip_deps//:platformdirs"],
+)

private/tests/override_module_repos/overridee/MODULE.bazel

@@ -0,0 +1,18 @@
+module(
+    name = "overridee",
+    version = "1.0",
+)
+
+bazel_dep(name = "rules_python", version = "0.34.0")
+bazel_dep(name = "rules_req_compile")
+local_path_override(
+    module_name = "rules_req_compile",
+    path = "../../../..",
+)
+
+requirements = use_extension("@rules_req_compile//extensions:python.bzl", "requirements")
+requirements.parse(
+    name = "pip_deps",
+    requirements_lock = ":requirements.txt",
+)
+use_repo(requirements, "pip_deps")