[Python Bazel] Allow overwriting grpc_library in python_rules.bzl (#35629)

Continues #35412 addressing feedback in https://github.com/grpc/grpc/pull/35412#issuecomment-1875980755

I'm unable to test the workspace with the newly added py_test, so I'm hoping a CI run here will tell me if the test works.

```console
~/grpc/test/distrib/bazel/python$ bazel --nohome_rc --nosystem_rc test //...
Loading:
Loading:
Loading: 0 packages loaded
Analyzing: 37 targets (0 packages loaded, 0 targets configured)
ERROR: /usr/local/foobar/home/sloretz/.cache/bazel/_bazel_sloretz/7f83b4f00f370e7c52a5cc586445673c/external/com_google_protobuf/upb_generator/BUILD:266:21: @com_google_protobuf//upb_generator:protoc-gen-upb_toolchain: no such attribute 'output_files' in 'proto_lang_toolchain' rule
ERROR: /usr/local/foobar/home/sloretz/.cache/bazel/_bazel_sloretz/7f83b4f00f370e7c52a5cc586445673c/external/com_google_protobuf/upb_generator/BUILD:305:21: @com_google_protobuf//upb_generator:protoc-gen-upb_minitable_toolchain: no such attribute 'output_files' in 'proto_lang_toolchain' rule
ERROR: /usr/local/foobar/home/sloretz/.cache/bazel/_bazel_sloretz/7f83b4f00f370e7c52a5cc586445673c/external/com_google_protobuf/upb_generator/BUILD:338:21: @com_google_protobuf//upb_generator:protoc-gen-upbdefs_toolchain: no such attribute 'output_files' in 'proto_lang_toolchain' rule
ERROR: /usr/local/foobar/home/sloretz/.cache/bazel/_bazel_sloretz/7f83b4f00f370e7c52a5cc586445673c/external/com_google_protobuf/upb_generator/BUILD:305:21: Target '@com_google_protobuf//upb_generator:protoc-gen-upb_minitable_stage1' contains an error and its package is in error and referenced by '@com_google_protobuf//upb_generator:protoc-gen-upb_minitable_toolchain'
ERROR: /usr/local/foobar/home/sloretz/.cache/bazel/_bazel_sloretz/7f83b4f00f370e7c52a5cc586445673c/external/com_github_grpc_grpc/src/proto/grpc/gcp/BUILD:19:14: every rule of type proto_library implicitly depends upon the target '@com_google_protobuf//upb_generator:protoc-gen-upb_minitable_toolchain', but this target could not be found because of: Target '@com_google_protobuf//upb_generator:protoc-gen-upb_minitable_toolchain' contains an error and its package is in error
ERROR: Analysis failed
ERROR: Analysis of target '//namespaced/upper/example:no_import_no_strip_py_pb2_grpc' failed; build aborted:
INFO: Elapsed time: 0.171s
INFO: 0 processes.
FAILED: Build did NOT complete successfully (0 packages loaded, 0 targets configured)
ERROR: Couldn't start the build. Unable to run tests
```

Closes #35629

PiperOrigin-RevId: 606695865

bazel/python_rules.bzl

@@ -222,11 +222,11 @@ def _generate_pb2_grpc_src_impl(context):
     py_info = _merge_pyinfos(
         [
             p,
-            context.attr._grpc_library[PyInfo],
+            context.attr.grpc_library[PyInfo],
         ] + [dep[PyInfo] for dep in context.attr.py_deps],
     )
 
-    runfiles = context.runfiles(files = out_files, transitive_files = py_info.transitive_sources).merge(context.attr._grpc_library[DefaultInfo].data_runfiles)
+    runfiles = context.runfiles(files = out_files, transitive_files = py_info.transitive_sources).merge(context.attr.grpc_library[DefaultInfo].data_runfiles)
 
     return [
         DefaultInfo(
@@ -261,7 +261,7 @@ _generate_pb2_grpc_src = rule(
             cfg = "exec",
             default = Label("//external:protocol_compiler"),
         ),
-        "_grpc_library": attr.label(
+        "grpc_library": attr.label(
             default = Label("//src/python/grpcio/grpc:grpcio"),
             providers = [PyInfo],
         ),
@@ -274,6 +274,7 @@ def py_grpc_library(
         srcs,
         deps,
         strip_prefixes = [],
+        grpc_library = Label("//src/python/grpcio/grpc:grpcio"),
         **kwargs):
     """Generate python code for gRPC services defined in a protobuf.
 
@@ -287,6 +288,9 @@ def py_grpc_library(
         stripped from the beginning of foo_pb2 modules imported by the
         generated stubs. This is useful in combination with the `imports`
         attribute of the `py_library` rule.
+      grpc_library: (`label`) a single `py_library` target representing the
+        python gRPC library target to be depended upon. This can be used to
+        generate code that depends on `grpcio` from the Python Package Index.
       **kwargs: Additional arguments to be supplied to the invocation of
         py_library.
     """
@@ -301,5 +305,6 @@ def py_grpc_library(
         deps = srcs,
         py_deps = deps,
         strip_prefixes = strip_prefixes,
+        grpc_library = grpc_library,
         **kwargs
     )

test/distrib/bazel/python/BUILD

@@ -170,3 +170,28 @@ py_test(
 python_rules_test_suite(
     name = "python_rules_test",
 )
+
+# Test that grpc_library attribute replaces grpcio dependency on
+# py_grpc_library targets
+
+py_library(
+    name = "grpc_library_replacement",
+    srcs = ["grpc_library_replacement.py"],
+)
+
+py_grpc_library(
+    name = "helloworld_py_pb2_grpc_library_changed",
+    srcs = [":helloworld_proto"],
+    grpc_library = ":grpc_library_replacement",
+    deps = [":helloworld_py_pb2"],
+)
+
+py_test(
+    name = "grpc_library_replacement_test",
+    srcs = ["grpc_library_replacement_test.py"],
+    main = "grpc_library_replacement_test.py",
+    python_version = "PY3",
+    deps = [
+        ":helloworld_py_pb2_grpc_library_changed",
+    ],
+)

test/distrib/bazel/python/grpc_library_replacement.py

@@ -0,0 +1,16 @@
+# Copyright 2024 gRPC authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Intentionally empty.
+# This is used to test the grpc_library attribute of py_grpc_library works.

test/distrib/bazel/python/grpc_library_replacement_test.py

@@ -0,0 +1,24 @@
+# Copyright 2024 gRPC authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+
+try:
+    import grpc
+except ImportError:
+    pass
+else:
+    sys.exit("Unexpectedly able to import grpc")
+
+import grpc_library_replacement