feat: allow registering arbitrary settings for py_binary transitions (#3248)

This implements the ability for users to add additional settings that
py_binary, py_test,
and py_wheel can transition on.

There were three main use cases motivating this feature:
1. Making it easier to have multiple pypi dependency closures and shared
dependencies.
2. Making it easier to override flags for `py_wheel`.
3. Making it easier to have per-target setting of things like
bootstrap_impl, venv
   site packages, etc.

It also adds most of our config settings to the the transition
inputs/outputs for those
rules, which allows users to per-target force particular settings
without having to
use e.g. `with_cfg` to wrap a target with the desired transition
settings. It also
lets use avoid adding dozens of attributes (one per setting); today
there are
about 17 flags.

Under the hood, this works by having a bzlmod api that users can pass
labels to. These
labels are put into a generated bzl file, which the rules load and add
to their
list of transition inputs/outputs. On the target level, the
`config_settings` attribute,
which is a `dict[label, str]`, can be set to change the particular flags
of interest.

Along the way...
* Create a common_labels.bzl file for the shared label strings
* Remove the defunct py_reconfig code in sh_py_run_test.

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: rickeylev

MODULE.bazel

@@ -13,9 +13,9 @@ bazel_dep(name = "platforms", version = "0.0.11")
 # Use py_proto_library directly from protobuf repository
 bazel_dep(name = "protobuf", version = "29.0-rc2", repo_name = "com_google_protobuf")
 
-internal_deps = use_extension("//python/private:internal_deps.bzl", "internal_deps")
+rules_python_config = use_extension("//python/extensions:config.bzl", "config")
 use_repo(
-    internal_deps,
+    rules_python_config,
     "pypi__build",
     "pypi__click",
     "pypi__colorama",
@@ -218,6 +218,19 @@ use_repo(
     "whl_with_build_files",
 )
 
+dev_rules_python_config = use_extension(
+    "//python/extensions:config.bzl",
+    "config",
+    dev_dependency = True,
+)
+dev_rules_python_config.add_transition_setting(
+    # Intentionally add a setting already present for testing
+    setting = "//python/config_settings:python_version",
+)
+dev_rules_python_config.add_transition_setting(
+    setting = "//tests/multi_pypi:external_deps_name",
+)
+
 # Add gazelle plugin so that we can run the gazelle example as an e2e integration
 # test and include the distribution files.
 local_path_override(
@@ -291,7 +304,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",
@@ -155,3 +157,26 @@ pip_parse(
 load("@dev_pip//:requirements.bzl", docs_install_deps = "install_deps")
 
 docs_install_deps()
+
+#####################
+# Pypi repos for //tests/multi_pypi
+
+pip_parse(
+    name = "pypi_alpha",
+    python_interpreter_target = interpreter,
+    requirements_lock = "//tests/multi_pypi/alpha:requirements.txt",
+)
+
+load("@pypi_alpha//:requirements.bzl", pypi_alpha_install_deps = "install_deps")
+
+pypi_alpha_install_deps()
+
+pip_parse(
+    name = "pypi_beta",
+    python_interpreter_target = interpreter,
+    requirements_lock = "//tests/multi_pypi/beta:requirements.txt",
+)
+
+load("@pypi_beta//:requirements.bzl", pypi_beta_install_deps = "install_deps")
+
+pypi_beta_install_deps()

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_beta",
+  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.
+
+```bzl
+# File: MODULE.bazel
+
+rules_python_config.add_transition_setting(
+    setting = "//:pypi_hub",
+)
+
+# File: BUILD.bazel
+
+```bzl
+
+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.

internal_dev_deps.bzl

@@ -41,7 +41,12 @@ def rules_python_internal_deps():
     For dependencies needed by *users* of rules_python, see
     python/private/py_repositories.bzl.
     """
-    internal_config_repo(name = "rules_python_internal")
+    internal_config_repo(
+        name = "rules_python_internal",
+        transition_settings = [
+            str(Label("//tests/multi_pypi:external_deps_name")),
+        ],
+    )
 
     local_repository(
         name = "other",

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:internal_config_repo_bzl",
+    ],
+)

python/extensions/config.bzl

@@ -0,0 +1,53 @@
+"""Extension for configuring global settings of rules_python."""
+
+load("//python/private:internal_config_repo.bzl", "internal_config_repo")
+load("//python/private/pypi:deps.bzl", "pypi_deps")
+
+_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."),
+    },
+)
+
+def _config_impl(mctx):
+    transition_setting_generators = {}
+    transition_settings = []
+    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_settings.append(setting)
+            transition_setting_generators[setting].append(mod.name)
+
+    internal_config_repo(
+        name = "rules_python_internal",
+        transition_setting_generators = transition_setting_generators,
+        transition_settings = transition_settings,
+    )
+
+    pypi_deps()
+
+config = module_extension(
+    doc = """Global settings for rules_python.
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
+""",
+    implementation = _config_impl,
+    tag_classes = {
+        "add_transition_setting": _add_transition_setting,
+    },
+)

python/private/BUILD.bazel

@@ -106,6 +106,7 @@ bzl_library(
     name = "builders_util_bzl",
     srcs = ["builders_util.bzl"],
     deps = [
+        ":bzlmod_enabled_bzl",
         "@bazel_skylib//lib:types",
     ],
 )
@@ -135,6 +136,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 +414,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 +590,7 @@ bzl_library(
     deps = [
         ":py_package_bzl",
         ":stamp_bzl",
+        ":transition_labels_bzl",
     ],
 )
 
@@ -649,6 +657,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_internal//: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_in_out_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_in_out_values("input", state[_INPUTS])
+    normalize_transition_in_out_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 = """\