fix: allowing to import code generated from proto with strip_import_prefix (#1406)

When the `proto_library` has `strip_import_prefix`, the py files from
proto are generated into a directory like
`bazel-bin/tests/py_proto_library/proto/_virtual_imports` and symlinked
in the runfiles as `<runfilesroot>/<workspace
name>/tests/py_proto_library/proto/_virtual_imports`.
We need to add `<workspace
name>/tests/py_proto_library/proto/_virtual_imports` to the `imports` of
`_PyProtoInfo`, so it will be appended to `PYTHONPATH`.

Modified an existing example to demonstrate the scenario and verify the
fix.

---------

Co-authored-by: Ignas Anikevicius <[email protected]>

Author: linzhp

.bazelrc

@@ -3,8 +3,8 @@
 # This lets us glob() up all the files inside the examples to make them inputs to tests
 # (Note, we cannot use `common --deleted_packages` because the bazel version command doesn't support it)
 # To update these lines, run tools/bazel_integration_test/update_deleted_packages.sh
-build --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_install,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
-query --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_install,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
+build --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_install,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/proto,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
+query --deleted_packages=examples/build_file_generation,examples/build_file_generation/random_number_generator,examples/bzlmod,examples/bzlmod_build_file_generation,examples/bzlmod_build_file_generation/other_module/other_module/pkg,examples/bzlmod_build_file_generation/runfiles,examples/bzlmod/entry_points,examples/bzlmod/entry_points/tests,examples/bzlmod/libs/my_lib,examples/bzlmod/other_module,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/bzlmod/tests,examples/bzlmod/tests/other_module,examples/bzlmod/whl_mods,examples/multi_python_versions/libs/my_lib,examples/multi_python_versions/requirements,examples/multi_python_versions/tests,examples/pip_install,examples/pip_parse,examples/pip_parse_vendored,examples/pip_repository_annotations,examples/py_proto_library,examples/py_proto_library/example.com/proto,tests/compile_pip_requirements,tests/compile_pip_requirements_test_from_external_workspace,tests/ignore_root_user_error,tests/pip_repository_entry_points
 
 test --test_output=errors
 

CHANGELOG.md

@@ -24,6 +24,7 @@ A brief description of the categories of changes:
 * Skip aliases for unloaded toolchains. Some Python versions that don't have full
   platform support, and referencing their undefined repositories can break operations
   like `bazel query rdeps(...)`.
+* Python code generated from `proto_library` with `strip_import_prefix` can be imported now.
 
 ## [0.26.0] - 2023-10-06
 

examples/py_proto_library/BUILD.bazel

@@ -1,22 +1,10 @@
-load("@rules_proto//proto:defs.bzl", "proto_library")
 load("@rules_python//python:defs.bzl", "py_test")
-load("@rules_python//python:proto.bzl", "py_proto_library")
-
-py_proto_library(
-    name = "pricetag_proto_py_pb2",
-    deps = [":pricetag_proto"],
-)
-
-proto_library(
-    name = "pricetag_proto",
-    srcs = ["pricetag.proto"],
-)
 
 py_test(
     name = "pricetag_test",
     srcs = ["test.py"],
     main = "test.py",
     deps = [
-        ":pricetag_proto_py_pb2",
+        "//example.com/proto:pricetag_proto_py_pb2",
     ],
 )

examples/py_proto_library/example.com/proto/BUILD.bazel

@@ -0,0 +1,15 @@
+load("@rules_proto//proto:defs.bzl", "proto_library")
+load("@rules_python//python:proto.bzl", "py_proto_library")
+
+py_proto_library(
+    name = "pricetag_proto_py_pb2",
+    visibility = ["//visibility:public"],
+    deps = [":pricetag_proto"],
+)
+
+proto_library(
+    name = "pricetag_proto",
+    srcs = ["pricetag.proto"],
+    # https://bazel.build/reference/be/protocol-buffer#proto_library.strip_import_prefix
+    strip_import_prefix = "/example.com",
+)

examples/py_proto_library/pricetag.proto renamed to examples/py_proto_library/example.com/proto/pricetag.proto

@@ -1,7 +1,7 @@
 import sys
 import unittest
 
-import pricetag_pb2
+from proto import pricetag_pb2
 
 
 class TestCase(unittest.TestCase):

examples/py_proto_library/test.py

@@ -66,6 +66,7 @@ def _py_proto_aspect_impl(target, ctx):
 
     generated_sources = []
     proto_info = target[ProtoInfo]
+    proto_root = proto_info.proto_source_root
     if proto_info.direct_sources:
         # Generate py files
         generated_sources = proto_common.declare_generated_files(
@@ -76,14 +77,11 @@ def _py_proto_aspect_impl(target, ctx):
         )
 
         # Handles multiple repository and virtual import cases
-        proto_root = proto_info.proto_source_root
         if proto_root.startswith(ctx.bin_dir.path):
-            plugin_output = proto_root
-        else:
-            plugin_output = ctx.bin_dir.path + "/" + proto_root
+            proto_root = proto_root[len(ctx.bin_dir.path) + 1:]
 
-        if plugin_output == ".":
-            plugin_output = ctx.bin_dir.path
+        plugin_output = ctx.bin_dir.path + "/" + proto_root
+        proto_root = ctx.workspace_name + "/" + proto_root
 
         proto_common.compile(
             actions = ctx.actions,
@@ -109,6 +107,10 @@ def _py_proto_aspect_impl(target, ctx):
     return [
         _PyProtoInfo(
             imports = depset(
+                # Adding to PYTHONPATH so the generated modules can be imported.
+                # This is necessary when there is strip_import_prefix, the Python
+                # modules are generated under _virtual_imports.
+                [proto_root],
                 transitive = [dep[PyInfo].imports for dep in api_deps],
             ),
             runfiles_from_proto_deps = runfiles_from_proto_deps,