tests: refactor py_reconfig rules so less boilerplate is needed to add attrs (#2933)

Just some minor refactoring of the py_reconfig rule so that it's easier
to add attributes
that affect transition state. After this, just two spots have to be
modified to
add an attribute (map of attrs, map of attr to transition label).

Author: rickeylev

tests/support/sh_py_run_test.bzl

@@ -13,14 +13,88 @@
 # limitations under the License.
 """Run a py_binary with altered config settings in an sh_test.
 
-This facilitates verify running binaries with different outer environmental
-settings and verifying their output without the overhead of a bazel-in-bazel
-integration test.
+This facilitates verify running binaries with different configuration settings
+without the overhead of a bazel-in-bazel integration test.
 """
 
 load("@rules_shell//shell:sh_test.bzl", "sh_test")
+load("//python/private:attr_builders.bzl", "attrb")  # buildifier: disable=bzl-visibility
+load("//python/private:py_binary_macro.bzl", "py_binary_macro")  # buildifier: disable=bzl-visibility
+load("//python/private:py_binary_rule.bzl", "create_py_binary_rule_builder")  # buildifier: disable=bzl-visibility
+load("//python/private:py_test_macro.bzl", "py_test_macro")  # buildifier: disable=bzl-visibility
+load("//python/private:py_test_rule.bzl", "create_py_test_rule_builder")  # buildifier: disable=bzl-visibility
 load("//python/private:toolchain_types.bzl", "TARGET_TOOLCHAIN_TYPE")  # buildifier: disable=bzl-visibility
-load(":py_reconfig.bzl", "py_reconfig_binary")
+load("//tests/support:support.bzl", "VISIBLE_FOR_TESTING")
+
+def _perform_transition_impl(input_settings, attr, base_impl):
+    settings = {k: input_settings[k] for k in _RECONFIG_INHERITED_OUTPUTS if k in input_settings}
+    settings.update(base_impl(input_settings, attr))
+
+    settings[VISIBLE_FOR_TESTING] = True
+    settings["//command_line_option:build_python_zip"] = attr.build_python_zip
+
+    for attr_name, setting_label in _RECONFIG_ATTR_SETTING_MAP.items():
+        if getattr(attr, attr_name):
+            settings[setting_label] = getattr(attr, attr_name)
+    return settings
+
+# Attributes that, if non-falsey (`if attr.<name>`), will copy their
+# value into the output settings
+_RECONFIG_ATTR_SETTING_MAP = {
+    "bootstrap_impl": "//python/config_settings:bootstrap_impl",
+    "extra_toolchains": "//command_line_option:extra_toolchains",
+    "python_src": "//python/bin:python_src",
+    "venvs_site_packages": "//python/config_settings:venvs_site_packages",
+    "venvs_use_declare_symlink": "//python/config_settings:venvs_use_declare_symlink",
+}
+
+_RECONFIG_INPUTS = _RECONFIG_ATTR_SETTING_MAP.values()
+_RECONFIG_OUTPUTS = _RECONFIG_INPUTS + [
+    "//command_line_option:build_python_zip",
+    VISIBLE_FOR_TESTING,
+]
+_RECONFIG_INHERITED_OUTPUTS = [v for v in _RECONFIG_OUTPUTS if v in _RECONFIG_INPUTS]
+
+_RECONFIG_ATTRS = {
+    "bootstrap_impl": attrb.String(),
+    "build_python_zip": attrb.String(default = "auto"),
+    "extra_toolchains": attrb.StringList(
+        doc = """
+Value for the --extra_toolchains flag.
+
+NOTE: You'll likely have to also specify //tests/support/cc_toolchains:all (or some CC toolchain)
+to make the RBE presubmits happy, which disable auto-detection of a CC
+toolchain.
+""",
+    ),
+    "python_src": attrb.Label(),
+    "venvs_site_packages": attrb.String(),
+    "venvs_use_declare_symlink": attrb.String(),
+}
+
+def _create_reconfig_rule(builder):
+    builder.attrs.update(_RECONFIG_ATTRS)
+
+    base_cfg_impl = builder.cfg.implementation()
+    builder.cfg.set_implementation(lambda *args: _perform_transition_impl(base_impl = base_cfg_impl, *args))
+    builder.cfg.update_inputs(_RECONFIG_INPUTS)
+    builder.cfg.update_outputs(_RECONFIG_OUTPUTS)
+    return builder.build()
+
+_py_reconfig_binary = _create_reconfig_rule(create_py_binary_rule_builder())
+
+_py_reconfig_test = _create_reconfig_rule(create_py_test_rule_builder())
+
+def py_reconfig_test(**kwargs):
+    """Create a py_test with customized build settings for testing.
+
+    Args:
+        **kwargs: kwargs to pass along to _py_reconfig_test.
+    """
+    py_test_macro(_py_reconfig_test, **kwargs)
+
+def py_reconfig_binary(**kwargs):
+    py_binary_macro(_py_reconfig_binary, **kwargs)
 
 def sh_py_run_test(*, name, sh_src, py_src, **kwargs):
     """Run a py_binary within a sh_test.