refactor: read migrated native flags through a centralized accessor function

python/config_settings/BUILD.bazel

@@ -1,4 +1,4 @@
-load("@bazel_skylib//rules:common_settings.bzl", "string_flag")
+load("@bazel_skylib//rules:common_settings.bzl", "bool_flag", "string_flag")
 load("@pythons_hub//:versions.bzl", "DEFAULT_PYTHON_VERSION", "MINOR_MAPPING", "PYTHON_VERSIONS")
 load(
     "//python/private:flags.bzl",
@@ -240,3 +240,9 @@ label_flag(
     # NOTE: Only public because it is used in pip hub repos.
     visibility = ["//visibility:public"],
 )
+
+bool_flag(
+    name = "experimental_python_import_all_repositories",
+    build_setting_default = True,
+    visibility = ["//visibility:public"],
+)

python/private/flags.bzl

@@ -21,6 +21,59 @@ unnecessary files when all that are needed are flag definitions.
 load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
 load(":enum.bzl", "FlagEnum", "enum")
 
+# Maps "--myflag" to a tuple of:
+#
+# - the flag's ctx.fragments native API accessor
+# -"native|starlark": which definition to use if the flag is available both
+#      from ctx.fragments and Starlark
+#
+# Builds that set --incompatible_remove_ctx_py_fragment or
+# --incompatible_remove_ctx_bazel_py_fragment disable ctx.fragments. These
+# builds assume flags are solely defined in Starlark.
+#
+# The "native|starlark" override is only for devs who are testing flag
+# Starlarkification. If ctx.fragments.[py|bazel_py] is available and
+# a flag is set to "starlark", we exclusively read its starlark version.
+#
+# See https://github.com/bazel-contrib/rules_python/issues/3252.
+_POSSIBLY_NATIVE_FLAGS = {
+    "build_python_zip": (lambda ctx: ctx.fragments.py.build_python_zip, "native"),
+    "default_to_explicit_init_py": (lambda ctx: ctx.fragments.py.default_to_explicit_init_py, "native"),
+    "disable_py2": (lambda ctx: ctx.fragments.py.disable_py2, "native"),
+    "python_import_all_repositories": (lambda ctx: ctx.fragments.bazel_py.python_import_all_repositories, "native"),
+    "python_path": (lambda ctx: ctx.fragments.bazel_py.python_path, "native"),
+}
+
+def read_possibly_native_flag(ctx, flag_name):
+    """
+    Canonical API for reading a Python build flag.
+
+    Flags might be defined in Starlark or native-Bazel. This function reasd flags
+    from tbe correct source based on supporting Bazel version and --incompatible*
+    flags that disable native references.
+
+    Args:
+        ctx: Rule's configuration context.
+        flag_name: Name of the flag to read, without preceding "--".
+
+    Returns:
+        The flag's value.
+    """
+
+    # Bazel 9.0+ can disable these fragments with --incompatible_remove_ctx_py_fragment and
+    # --incompatible_remove_ctx_bazel_py_fragment. Disabling them means bazel expects
+    # Python to read Starlark flags.
+    use_native_def = hasattr(ctx.fragments, "py") and hasattr(ctx.fragments, "bazel_py")
+
+    # Developer override to force the Starlark definition for testing.
+    if _POSSIBLY_NATIVE_FLAGS[flag_name][1] == "starlark":
+        use_native_def = False
+    if use_native_def:
+        return _POSSIBLY_NATIVE_FLAGS[flag_name][0](ctx)
+    else:
+        # Starlark definition of "--foo" is assumed to be a label dependency named "_foo".
+        return getattr(ctx.attr, "_" + flag_name)[BuildSettingInfo].value
+
 def _AddSrcsToRunfilesFlag_is_enabled(ctx):
     value = ctx.attr._add_srcs_to_runfiles_flag[BuildSettingInfo].value
     if value == AddSrcsToRunfilesFlag.AUTO:

python/private/py_executable.bzl

@@ -53,7 +53,7 @@ load(
     "target_platform_has_any_constraint",
 )
 load(":common_labels.bzl", "labels")
-load(":flags.bzl", "BootstrapImplFlag", "VenvsUseDeclareSymlinkFlag")
+load(":flags.bzl", "BootstrapImplFlag", "VenvsUseDeclareSymlinkFlag", "read_possibly_native_flag")
 load(":precompile.bzl", "maybe_precompile")
 load(":py_cc_link_params_info.bzl", "PyCcLinkParamsInfo")
 load(":py_executable_info.bzl", "PyExecutableInfo")
@@ -293,7 +293,7 @@ def _get_stamp_flag(ctx):
 
 def _should_create_init_files(ctx):
     if ctx.attr.legacy_create_init == -1:
-        return not ctx.fragments.py.default_to_explicit_init_py
+        return not read_possibly_native_flag(ctx, "default_to_explicit_init_py")
     else:
         return bool(ctx.attr.legacy_create_init)
 
@@ -381,7 +381,7 @@ def _create_executable(
     extra_files_to_build = []
 
     # NOTE: --build_python_zip defaults to true on Windows
-    build_zip_enabled = ctx.fragments.py.build_python_zip
+    build_zip_enabled = read_possibly_native_flag(ctx, "build_python_zip")
 
     # When --build_python_zip is enabled, then the zip file becomes
     # one of the default outputs.
@@ -587,7 +587,7 @@ def _create_venv(ctx, output_prefix, imports, runtime_details):
         output = site_init,
         substitutions = {
             "%coverage_tool%": _get_coverage_tool_runfiles_path(ctx, runtime),
-            "%import_all%": "True" if ctx.fragments.bazel_py.python_import_all_repositories else "False",
+            "%import_all%": "True" if read_possibly_native_flag(ctx, "python_import_all_repositories") else "False",
             "%site_init_runfiles_path%": "{}/{}".format(ctx.workspace_name, site_init.short_path),
             "%workspace_name%": ctx.workspace_name,
         },
@@ -668,7 +668,7 @@ def _create_stage2_bootstrap(
         output = output,
         substitutions = {
             "%coverage_tool%": _get_coverage_tool_runfiles_path(ctx, runtime),
-            "%import_all%": "True" if ctx.fragments.bazel_py.python_import_all_repositories else "False",
+            "%import_all%": "True" if read_possibly_native_flag(ctx, "python_import_all_repositories") else "False",
             "%imports%": ":".join(imports.to_list()),
             "%main%": main_py_path,
             "%main_module%": ctx.attr.main_module,
@@ -755,7 +755,7 @@ def _create_stage1_bootstrap(
             template = ctx.file._bootstrap_template
 
         subs["%coverage_tool%"] = coverage_tool_runfiles_path
-        subs["%import_all%"] = ("True" if ctx.fragments.bazel_py.python_import_all_repositories else "False")
+        subs["%import_all%"] = ("True" if read_possibly_native_flag(ctx, "python_import_all_repositories") else "False")
         subs["%imports%"] = ":".join(imports.to_list())
         subs["%main%"] = "{}/{}".format(ctx.workspace_name, main_py.short_path)
 
@@ -1135,7 +1135,7 @@ def _get_runtime_details(ctx, semantics):
     #
     # TOOD(bazelbuild/bazel#7901): Remove this once --python_path flag is removed.
 
-    flag_interpreter_path = ctx.fragments.bazel_py.python_path
+    flag_interpreter_path = read_possibly_native_flag(ctx, "python_path")
     toolchain_runtime, effective_runtime = _maybe_get_runtime_from_ctx(ctx)
     if not effective_runtime:
         # Clear these just in case

python/private/py_runtime_pair_rule.bzl

@@ -17,6 +17,7 @@
 load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
 load("//python:py_runtime_info.bzl", "PyRuntimeInfo")
 load(":common_labels.bzl", "labels")
+load(":flags.bzl", "read_possibly_native_flag")
 load(":reexports.bzl", "BuiltinPyRuntimeInfo")
 load(":util.bzl", "IS_BAZEL_7_OR_HIGHER")
 
@@ -69,7 +70,7 @@ def _is_py2_disabled(ctx):
     # TODO: Remove this once all supported Balze versions have this flag.
     if not hasattr(ctx.fragments.py, "disable_py"):
         return False
-    return ctx.fragments.py.disable_py2
+    return read_possibly_native_flag(ctx, "disable_py2")
 
 _MaybeBuiltinPyRuntimeInfo = [[BuiltinPyRuntimeInfo]] if BuiltinPyRuntimeInfo != None else []