fix(bazel/python): incompatible_enable_proto_toolchain_resolution (#17078)

This was originally landed in https://github.com/bazelbuild/rules_python/pull/1577 However the fork of py_proto_library brought to the protobuf repo was from an earlier commit: 8257c4469d (cl/623401031 for googlers)
Thus the fix was lost. This PR simply cherry-picks that change.

Closes #17078

COPYBARA_INTEGRATE_REVIEW=https://github.com/protocolbuffers/protobuf/pull/17078 from alexeagle:py_proto_toolchain 226540ef3a
PiperOrigin-RevId: 642123868

bazel/py_proto_library.bzl

@@ -4,7 +4,7 @@ load("@rules_python//python:py_info.bzl", "PyInfo")
 load("//bazel/common:proto_common.bzl", "proto_common")
 load("//bazel/common:proto_info.bzl", "ProtoInfo")
 
-ProtoLangToolchainInfo = proto_common.ProtoLangToolchainInfo
+PY_PROTO_TOOLCHAIN = "@rules_python//python/proto:toolchain_type"
 
 _PyProtoInfo = provider(
     doc = "Encapsulates information needed by the Python proto rules.",
@@ -22,6 +22,9 @@ _PyProtoInfo = provider(
 def _filter_provider(provider, *attrs):
     return [dep[provider] for attr in attrs for dep in attr if provider in dep]
 
+def _incompatible_toolchains_enabled():
+    return getattr(proto_common, "INCOMPATIBLE_ENABLE_PROTO_TOOLCHAIN_RESOLUTION", False)
+
 def _py_proto_aspect_impl(target, ctx):
     """Generates and compiles Python code for a proto_library.
 
@@ -48,7 +51,14 @@ def _py_proto_aspect_impl(target, ctx):
                 proto.path,
             ))
 
-    proto_lang_toolchain_info = ctx.attr._aspect_proto_toolchain[ProtoLangToolchainInfo]
+    if _incompatible_toolchains_enabled():
+        toolchain = ctx.toolchains[PY_PROTO_TOOLCHAIN]
+        if not toolchain:
+            fail("No toolchains registered for '%s'." % PY_PROTO_TOOLCHAIN)
+        proto_lang_toolchain_info = toolchain.proto
+    else:
+        proto_lang_toolchain_info = getattr(ctx.attr, "_aspect_proto_toolchain")[proto_common.ProtoLangToolchainInfo]
+
     api_deps = [proto_lang_toolchain_info.runtime]
 
     generated_sources = []
@@ -110,7 +120,7 @@ def _py_proto_aspect_impl(target, ctx):
 
 _py_proto_aspect = aspect(
     implementation = _py_proto_aspect_impl,
-    attrs = {
+    attrs = {} if _incompatible_toolchains_enabled() else {
         "_aspect_proto_toolchain": attr.label(
             default = "//python:python_toolchain",
         ),
@@ -118,6 +128,7 @@ _py_proto_aspect = aspect(
     attr_aspects = ["deps"],
     required_providers = [ProtoInfo],
     provides = [_PyProtoInfo],
+    toolchains = [PY_PROTO_TOOLCHAIN] if _incompatible_toolchains_enabled() else [],
 )
 
 def _py_proto_library_rule(ctx):