From c8f467aea99e074776901fc4c53d59565b90d988 Mon Sep 17 00:00:00 2001
From: Ming Lyu <CareF.Lm@gmail.com>
Date: Thu, 17 Aug 2023 05:23:23 +0800
Subject: [PATCH] [Bazel] Enable grpcio-reflection to be used via Bazel
 (#31013)

The current `py_grpc_library` results in the wrong grpc proto python
code path when grpc is a third-party source code in a Bazel project.
This PR should fix it.

fixes #31011

<!--

If you know who should review your pull request, please assign it to
that
person, otherwise the pull request would get assigned randomly.

If your pull request is for a specific language, please add the
appropriate
lang label.

-->

---------

Co-authored-by: Richard Belleville <rbellevi@google.com>
---
 bazel/python_rules.bzl                          |  9 +++++++--
 test/distrib/bazel/python/BUILD                 | 10 ++++++++++
 .../python/import_from_grpcio_reflection.py     | 17 +++++++++++++++++
 3 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 test/distrib/bazel/python/import_from_grpcio_reflection.py

diff --git a/bazel/python_rules.bzl b/bazel/python_rules.bzl
index 648af98e68c..88af99757eb 100644
--- a/bazel/python_rules.bzl
+++ b/bazel/python_rules.bzl
@@ -140,7 +140,7 @@ def _generate_py_impl(context):
         for py_src in context.attr.deps[0][PyProtoInfo].generated_py_srcs:
             reimport_py_file = context.actions.declare_file(py_src.basename)
             py_sources.append(reimport_py_file)
-            import_line = "from %s import *" % py_src.short_path.replace("/", ".")[:-len(".py")]
+            import_line = "from %s import *" % py_src.short_path.replace("..", "external").replace("/", ".")[:-len(".py")]
             context.actions.write(reimport_py_file, import_line)
 
     # Collect output PyInfo provider.
@@ -190,10 +190,15 @@ def _generate_pb2_grpc_src_impl(context):
     arguments = []
     tools = [context.executable._protoc, context.executable._grpc_plugin]
     out_dir = get_out_dir(protos, context)
+    if out_dir.import_path:
+        # is virtual imports
+        out_path = out_dir.path
+    else:
+        out_path = context.genfiles_dir.path
     arguments += get_plugin_args(
         context.executable._grpc_plugin,
         plugin_flags,
-        out_dir.path,
+        out_path,
         False,
     )
 
diff --git a/test/distrib/bazel/python/BUILD b/test/distrib/bazel/python/BUILD
index 74b69249d02..f784dd8796f 100644
--- a/test/distrib/bazel/python/BUILD
+++ b/test/distrib/bazel/python/BUILD
@@ -145,6 +145,16 @@ py_test(
     ],
 )
 
+py_test(
+    name = "import_from_grpcio_reflection_test",
+    srcs = ["import_from_grpcio_reflection.py"],
+    main = "import_from_grpcio_reflection.py",
+    python_version = "PY3",
+    deps = [
+        "@com_github_grpc_grpc//src/python/grpcio_reflection/grpc_reflection/v1alpha:grpc_reflection",
+    ],
+)
+
 # Test that a py_proto_library can be successfully imported without requiring
 # explicit dependencies on unused dependencies of the generated code.
 py_test(
diff --git a/test/distrib/bazel/python/import_from_grpcio_reflection.py b/test/distrib/bazel/python/import_from_grpcio_reflection.py
new file mode 100644
index 00000000000..876c2853a45
--- /dev/null
+++ b/test/distrib/bazel/python/import_from_grpcio_reflection.py
@@ -0,0 +1,17 @@
+# Copyright 2022 the 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.
+"""A trivial executable that imports grpc_reflection.
+"""
+
+from grpc_reflection.v1alpha import reflection