feat: allow registering arbitrary settings for py_binary transitions

Author: rickeylev

MODULE.bazel

@@ -34,6 +34,20 @@ use_repo(
     "rules_python_internal",
 )
 
+rules_python_config = use_extension(
+    "//python/extensions:config.bzl",
+    "rules_python_config",
+    dev_dependency = True,
+)
+use_repo(rules_python_config, "rules_python_config")
+rules_python_config.add_transition_setting(
+    # Intentionally add a setting already present for testing
+    setting = "//python/config_settings:python_version",
+)
+rules_python_config.add_transition_setting(
+    setting = "//tests/multi_pypi:external_deps_name",
+)
+
 # We need to do another use_extension call to expose the "pythons_hub"
 # repo.
 python = use_extension("//python/extensions:python.bzl", "python")
@@ -291,7 +305,17 @@ dev_pip.parse(
     python_version = "3.11",
     requirements_lock = "//examples/wheel:requirements_server.txt",
 )
-use_repo(dev_pip, "dev_pip", "pypiserver")
+dev_pip.parse(
+    hub_name = "pypi_alpha",
+    python_version = "3.11",
+    requirements_lock = "//tests/multi_pypi/alpha:requirements.txt",
+)
+dev_pip.parse(
+    hub_name = "pypi_beta",
+    python_version = "3.11",
+    requirements_lock = "//tests/multi_pypi/beta:requirements.txt",
+)
+use_repo(dev_pip, "dev_pip", "pypi_alpha", "pypi_beta", "pypiserver")
 
 # Bazel integration test setup below
 

WORKSPACE

@@ -69,7 +69,9 @@ load("//:internal_dev_setup.bzl", "rules_python_internal_setup")
 rules_python_internal_setup()
 
 load("@pythons_hub//:versions.bzl", "PYTHON_VERSIONS")
-load("//python:repositories.bzl", "python_register_multi_toolchains")
+load("//python:repositories.bzl", "py_repositories", "python_register_multi_toolchains")
+
+py_repositories()
 
 python_register_multi_toolchains(
     name = "python",

docs/howto/common-deps-with-multipe-pypi-versions.md

@@ -0,0 +1,91 @@
+# How to use a common set of dependencies with multiple PyPI versions
+
+In this guide, we show how to handle a situation common to monorepos
+that extensively share code: How does a common library refer to the correct
+`@pypi_<name>` hub when binaries may have their own requirements (and thus
+PyPI hub name)? Stated as code, this situation:
+
+```bzl
+
+py_binary(
+  name = "bin_alpha"
+  deps = ["@pypi_alpha//requests", ":common"],
+)
+py_binary(
+  name = "bin_alpha"
+  deps = ["@pypi_beta//requests", ":common"],
+)
+
+py_library(
+  name = "common",
+  deps = ["@pypi_???//more_itertools"] # <-- Which @pypi repo?
+)
+```
+
+## Using flags to pick a hub
+
+The basic trick to make `:common` pick the appropriate `@pypi_<name>` is to use
+`select()` to choose one based on build flags. To help this process, `py_binary`
+et al allow forcing particular build flags to be used, and custom flags can be
+registered to allow `py_binary` et al to set them.
+
+In this example, we create a custom string flag named `//:pypi_hub`,
+register it to allow using it with `py_binary` directly, then use `select()`
+to pick different dependencies.
+
+```
+# File: MODULE.bazel
+
+rules_python_config.add_transition_setting(
+    setting = "@//pick:pypi_hub",
+)
+
+# File: BUILD.bazel
+
+```bazel
+
+load("@bazel_skylib//rules:common_settings.bzl", "string_flag")
+
+string_flag(
+    name = "pypi_hub",
+)
+
+config_setting(
+    name = "is_pypi_alpha",
+    flag_values = {"//:pypi_hub": "alpha"},
+)
+
+config_setting(
+    name = "is_pypi_beta",
+    flag_values = {"//:pypi_hub": "beta"}
+)
+
+py_binary(
+    name = "bin_alpha",
+    srcs = ["bin_alpha.py"],
+    config_settings = {
+        "//:pypi_hub": "alpha",
+    },
+    deps ["@pypi_alpha//requests", ":common"],
+)
+py_binary(
+    name = "bin_beta",
+    srcs = ["bin_beta.py"],
+    config_settings = {
+        "//:pypi_hub": "beta",
+    },
+    deps ["@pypi_beta//requests", ":common"],
+)
+py_library(
+    name = "common",
+    deps = select({
+        ":is_pypi_alpha": ["@pypi_alpha//more_itertools"],
+        ":is_pypi_beta": ["@pypi_beta//more_itertools"],
+    }),
+)
+```
+
+When `bin_alpha` and `bin_beta` are built, they will have the `pypi_hub`
+flag force to their respective value. When `:common` is evaluated, it sees
+the flag value of the binary that is consuming it, and the `select()` resolves
+appropriately.

python/extensions/BUILD.bazel

@@ -39,3 +39,12 @@ bzl_library(
         "//python/private:python_bzl",
     ],
 )
+
+bzl_library(
+    name = "config_bzl",
+    srcs = ["config.bzl"],
+    visibility = ["//:__subpackages__"],
+    deps = [
+        "//python/private:rules_python_config_repo_bzl",
+    ],
+)

python/extensions/config.bzl

@@ -0,0 +1,48 @@
+"""Extension for configuring global settings of rules_python."""
+
+load("//python/private:rules_python_config_repo.bzl", "rules_python_config_repo")
+
+def _rules_python_config_impl(mctx):
+    transition_setting_generators = {}
+    for mod in mctx.modules:
+        for tag in mod.tags.add_transition_setting:
+            setting = str(tag.setting)
+            if setting not in transition_setting_generators:
+                transition_setting_generators[setting] = []
+            transition_setting_generators[setting].append(mod.name)
+
+    rules_python_config_repo(
+        name = "rules_python_config",
+        transition_setting_generators = transition_setting_generators,
+        transition_settings = transition_setting_generators.keys(),
+    )
+
+_add_transition_setting = tag_class(
+    doc = """
+Specify a build setting that terminal rules transition on by default.
+
+Terminal rules are rules such as py_binary, py_test, py_wheel, or similar
+rules that represent some deployable unit. Settings added here can
+then be used a keys with the {obj}`config_settings` attribute.
+
+:::{note}
+This adds the label as a dependency of the Python rules. Take care to not refer
+to repositories that are expensive to create or invalidate frequently.
+:::
+""",
+    attrs = {
+        "setting": attr.label(doc = "The build setting to add."),
+    },
+)
+
+rules_python_config = module_extension(
+    doc = """Global settings for rules_python.
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
+""",
+    implementation = _rules_python_config_impl,
+    tag_classes = {
+        "add_transition_setting": _add_transition_setting,
+    },
+)

python/private/BUILD.bazel

@@ -135,6 +135,11 @@ bzl_library(
     ],
 )
 
+bzl_library(
+    name = "common_labels_bzl",
+    srcs = ["common_labels.bzl"],
+)
+
 bzl_library(
     name = "config_settings_bzl",
     srcs = ["config_settings.bzl"],
@@ -408,6 +413,7 @@ bzl_library(
         ":py_runtime_info_bzl",
         ":rules_cc_srcs_bzl",
         ":toolchain_types_bzl",
+        ":transition_labels_bzl",
         "@bazel_skylib//lib:dicts",
         "@bazel_skylib//lib:paths",
         "@bazel_skylib//lib:structs",
@@ -583,6 +589,7 @@ bzl_library(
     deps = [
         ":py_package_bzl",
         ":stamp_bzl",
+        ":transition_labels_bzl",
     ],
 )
 
@@ -618,6 +625,14 @@ bzl_library(
     ],
 )
 
+bzl_library(
+    name = "rules_python_config_repo_bzl",
+    srcs = ["rules_python_config_repo.bzl"],
+    deps = [
+        ":repo_utils_bzl",
+    ],
+)
+
 bzl_library(
     name = "sentinel_bzl",
     srcs = ["sentinel.bzl"],
@@ -649,6 +664,16 @@ bzl_library(
     srcs = ["toolchain_types.bzl"],
 )
 
+bzl_library(
+    name = "transition_labels_bzl",
+    srcs = ["transition_labels.bzl"],
+    deps = [
+        "common_labels_bzl",
+        "@bazel_skylib//lib:collections",
+        "@rules_python_config//:extra_transition_settings_bzl",
+    ],
+)
+
 bzl_library(
     name = "util_bzl",
     srcs = ["util.bzl"],

python/private/attr_builders.bzl

@@ -31,6 +31,7 @@ load(
     "kwargs_setter",
     "kwargs_setter_doc",
     "kwargs_setter_mandatory",
+    "normalize_transition_inout_values",
     "to_label_maybe",
 )
 
@@ -167,6 +168,8 @@ def _AttrCfg_new(
     }
     kwargs_set_default_list(state, _INPUTS)
     kwargs_set_default_list(state, _OUTPUTS)
+    normalize_transition_inout_values("input", state[_INPUTS])
+    normalize_transition_inout_values("output", state[_OUTPUTS])
 
     # buildifier: disable=uninitialized
     self = struct(

python/private/attributes.bzl

@@ -405,8 +405,58 @@ COVERAGE_ATTRS = {
 # Attributes specific to Python executable-equivalent rules. Such rules may not
 # accept Python sources (e.g. some packaged-version of a py_test/py_binary), but
 # still accept Python source-agnostic settings.
+CONFIG_SETTINGS_ATTR = {
+    "config_settings": lambda: attrb.LabelKeyedStringDict(
+        doc = """
+Config settings to change for this target.
+
+The keys are labels for settings, and the values are strings for the new value
+to use. Pass `Label` objects or canonical label strings for the keys to ensure
+they resolve as expected (canonical labels start with `@@` and can be
+obtained by calling `str(Label(...))`).
+
+Most `@rules_python//python/config_setting` settings can be used here, which
+allows, for example, making only a certain `py_binary` use
+{obj}`--boostrap_impl=script`.
+
+Additional or custom config settings can be registered using the
+{obj}`add_transition_setting` API. This allows, for example, forcing a
+particular CPU, or defining a custom setting that `select()` uses elsewhere
+to pick between `pip.parse` hubs. See the [How to guide on multiple
+versions of a library] for a more concrete example.
+
+:::{note}
+These values are transitioned on, so will affect the analysis graph and the
+associated memory overhead. The more unique configurations in your overall
+build, the more memory and (often unnecessary) re-analysis and re-building
+can occur. See
+https://bazel.build/extending/config#memory-performance-considerations for
+more information about risks and considerations.
+:::
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
+""",
+    ),
+}
+
+def apply_config_settings_attr(settings, attr):
+    """Applies the config_settings attribute to the settings.
+
+    Args:
+        settings: The settings dict to modify in-place.
+        attr: The rule attributes struct.
+
+    Returns:
+        {type}`dict[str, object]` the input `settings` value.
+    """
+    for key, value in attr.config_settings.items():
+        settings[str(key)] = value
+    return settings
+
 AGNOSTIC_EXECUTABLE_ATTRS = dicts.add(
     DATA_ATTRS,
+    CONFIG_SETTINGS_ATTR,
     {
         "env": lambda: attrb.StringDict(
             doc = """\