fix: make py_proto_library respect PyInfo imports (#1046)

py_proto_library has an implicitly dependency on the protobuf client runtime, and it was ending up in runfiles, but it wasn't imported because the `imports` value it was setting wasn't be propagated.

To fix, make py_proto_library properly propagate the imports attribute the protobuf client runtime so that the libraries are added to sys.path correct.

Also adds an example for bzlmod and old way of using the py_proto_library as a test.

---------

Co-authored-by: Richard Levasseur <richardlev@gmail.com>
diff --git a/.bazelci/presubmit.yml b/.bazelci/presubmit.yml
index 76f9d8b..706c655 100644
--- a/.bazelci/presubmit.yml
+++ b/.bazelci/presubmit.yml
@@ -31,6 +31,11 @@
     - "..."
   test_flags:
     - "--test_tag_filters=-integration-test"
+.common_bzlmod_flags: &common_bzlmod_flags
+  test_flags:
+    - "--experimental_enable_bzlmod"
+  build_flags:
+    - "--experimental_enable_bzlmod"
 .reusable_build_test_all: &reusable_build_test_all
   build_targets: ["..."]
   test_targets: ["..."]
@@ -211,6 +216,53 @@
   # We don't run pip_parse_vendored under Windows as the file checked in is
   # generated from a repository rule containing OS-specific rendered paths.
 
+  integration_test_py_proto_library_ubuntu:
+    <<: *reusable_build_test_all
+    name: py_proto_library integration tests on Ubuntu
+    working_directory: examples/py_proto_library
+    platform: ubuntu2004
+  integration_test_py_proto_library_debian:
+    <<: *reusable_build_test_all
+    name: py_proto_library integration tests on Debian
+    working_directory: examples/py_proto_library
+    platform: debian11
+  integration_test_py_proto_library_macos:
+    <<: *reusable_build_test_all
+    name: py_proto_library integration tests on macOS
+    working_directory: examples/py_proto_library
+    platform: macos
+  integration_test_py_proto_library_windows:
+    <<: *reusable_build_test_all
+    name: py_proto_library integration tests on Windows
+    working_directory: examples/py_proto_library
+    platform: windows
+
+  # Check the same using bzlmod as well
+  integration_test_py_proto_library_bzlmod_ubuntu:
+    <<: *reusable_build_test_all
+    <<: *common_bzlmod_flags
+    name: py_proto_library bzlmod integration tests on Ubuntu
+    working_directory: examples/py_proto_library
+    platform: ubuntu2004
+  integration_test_py_proto_library_bzlmod_debian:
+    <<: *reusable_build_test_all
+    <<: *common_bzlmod_flags
+    name: py_proto_library bzlmod integration tests on Debian
+    working_directory: examples/py_proto_library
+    platform: debian11
+  integration_test_py_proto_library_bzlmod_macos:
+    <<: *reusable_build_test_all
+    <<: *common_bzlmod_flags
+    name: py_proto_library bzlmod integration tests on macOS
+    working_directory: examples/py_proto_library
+    platform: macos
+  integration_test_py_proto_library_bzlmod_windows:
+    <<: *reusable_build_test_all
+    <<: *common_bzlmod_flags
+    name: py_proto_library bzlmod integration tests on Windows
+    working_directory: examples/py_proto_library
+    platform: windows
+
   integration_test_pip_repository_annotations_ubuntu:
     <<: *reusable_build_test_all
     name: pip_repository_annotations integration tests on Ubuntu
diff --git a/.bazelrc b/.bazelrc
index 2dc3259..d607cdd 100644
--- a/.bazelrc
+++ b/.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/get_url,examples/bzlmod,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/multi_python_versions,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_import,examples/relative_requirements,tests/compile_pip_requirements,tests/pip_repository_entry_points,tests/pip_deps
-query --deleted_packages=examples/build_file_generation,examples/build_file_generation/get_url,examples/bzlmod,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/multi_python_versions,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_import,examples/relative_requirements,tests/compile_pip_requirements,tests/pip_repository_entry_points,tests/pip_deps
+build --deleted_packages=examples/build_file_generation,examples/build_file_generation/get_url,examples/bzlmod,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/multi_python_versions,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_import,examples/py_proto_library,examples/relative_requirements,tests/compile_pip_requirements,tests/pip_repository_entry_points,tests/pip_deps
+query --deleted_packages=examples/build_file_generation,examples/build_file_generation/get_url,examples/bzlmod,examples/bzlmod/other_module/other_module/pkg,examples/bzlmod/runfiles,examples/multi_python_versions,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_import,examples/py_proto_library,examples/relative_requirements,tests/compile_pip_requirements,tests/pip_repository_entry_points,tests/pip_deps
 
 test --test_output=errors
 
diff --git a/examples/BUILD.bazel b/examples/BUILD.bazel
index e0a7e5a..3ef8905 100644
--- a/examples/BUILD.bazel
+++ b/examples/BUILD.bazel
@@ -44,6 +44,18 @@
 )
 
 bazel_integration_test(
+    name = "py_proto_library_example",
+    timeout = "long",
+)
+
+bazel_integration_test(
+    name = "py_proto_library_example_bzlmod",
+    timeout = "long",
+    bzlmod = True,
+    dirname = "py_proto_library",
+)
+
+bazel_integration_test(
     name = "multi_python_versions_example",
     timeout = "long",
 )
diff --git a/examples/py_proto_library/.bazelrc b/examples/py_proto_library/.bazelrc
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/examples/py_proto_library/.bazelrc
diff --git a/examples/py_proto_library/.gitignore b/examples/py_proto_library/.gitignore
new file mode 100644
index 0000000..e5ae073
--- /dev/null
+++ b/examples/py_proto_library/.gitignore
@@ -0,0 +1,4 @@
+# git ignore patterns
+
+/bazel-*
+user.bazelrc
diff --git a/examples/py_proto_library/BUILD.bazel b/examples/py_proto_library/BUILD.bazel
new file mode 100644
index 0000000..7a18a5e
--- /dev/null
+++ b/examples/py_proto_library/BUILD.bazel
@@ -0,0 +1,22 @@
+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",
+    ],
+)
diff --git a/examples/py_proto_library/MODULE.bazel b/examples/py_proto_library/MODULE.bazel
new file mode 100644
index 0000000..5ce0924
--- /dev/null
+++ b/examples/py_proto_library/MODULE.bazel
@@ -0,0 +1,28 @@
+module(
+    name = "rules_python_py_proto_library_example",
+    version = "0.0.0",
+    compatibility_level = 1,
+)
+
+bazel_dep(name = "rules_python", version = "0.17.3")
+
+# The following local_path_override is only needed to run this example as part of our CI.
+local_path_override(
+    module_name = "rules_python",
+    path = "../..",
+)
+
+python = use_extension("@rules_python//python:extensions.bzl", "python")
+python.toolchain(
+    name = "python3_9",
+    configure_coverage_tool = True,
+    python_version = "3.9",
+)
+use_repo(python, "python3_9_toolchains")
+
+register_toolchains(
+    "@python3_9_toolchains//:all",
+)
+
+# We are using rules_proto to define rules_proto targets to be consumed by py_proto_library.
+bazel_dep(name = "rules_proto", version = "5.3.0-21.7")
diff --git a/examples/py_proto_library/WORKSPACE b/examples/py_proto_library/WORKSPACE
new file mode 100644
index 0000000..bf38112
--- /dev/null
+++ b/examples/py_proto_library/WORKSPACE
@@ -0,0 +1,48 @@
+workspace(name = "rules_python_py_proto_library_example")
+
+# The following local_path_override is only needed to run this example as part of our CI.
+local_repository(
+    name = "rules_python",
+    path = "../..",
+)
+
+# When not using this example in the rules_python git repo you would load the python
+# rules using http_archive(), as documented in the release notes.
+
+load("@rules_python//python:repositories.bzl", "py_repositories", "python_register_toolchains")
+
+# We install the rules_python dependencies using the function below.
+py_repositories()
+
+python_register_toolchains(
+    name = "python39",
+    python_version = "3.9",
+)
+
+# Then we need to setup dependencies in order to use py_proto_library
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+http_archive(
+    name = "rules_proto",
+    sha256 = "dc3fb206a2cb3441b485eb1e423165b231235a1ea9b031b4433cf7bc1fa460dd",
+    strip_prefix = "rules_proto-5.3.0-21.7",
+    urls = [
+        "https://github.com/bazelbuild/rules_proto/archive/refs/tags/5.3.0-21.7.tar.gz",
+    ],
+)
+
+http_archive(
+    name = "com_google_protobuf",
+    sha256 = "75be42bd736f4df6d702a0e4e4d30de9ee40eac024c4b845d17ae4cc831fe4ae",
+    strip_prefix = "protobuf-21.7",
+    urls = [
+        "https://mirror.bazel.build/github.com/protocolbuffers/protobuf/archive/v21.7.tar.gz",
+        "https://github.com/protocolbuffers/protobuf/archive/v21.7.tar.gz",
+    ],
+)
+
+load("@rules_proto//proto:repositories.bzl", "rules_proto_dependencies", "rules_proto_toolchains")
+
+rules_proto_dependencies()
+
+rules_proto_toolchains()
diff --git a/examples/py_proto_library/WORKSPACE.bzlmod b/examples/py_proto_library/WORKSPACE.bzlmod
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/examples/py_proto_library/WORKSPACE.bzlmod
diff --git a/examples/py_proto_library/pricetag.proto b/examples/py_proto_library/pricetag.proto
new file mode 100644
index 0000000..c952248
--- /dev/null
+++ b/examples/py_proto_library/pricetag.proto
@@ -0,0 +1,8 @@
+syntax = "proto3";
+
+package rules_python;
+
+message PriceTag {
+  string name = 2;
+  double cost = 1;
+}
diff --git a/examples/py_proto_library/test.py b/examples/py_proto_library/test.py
new file mode 100644
index 0000000..9f09702
--- /dev/null
+++ b/examples/py_proto_library/test.py
@@ -0,0 +1,17 @@
+import sys
+import unittest
+
+import pricetag_pb2
+
+
+class TestCase(unittest.TestCase):
+    def test_pricetag(self):
+        got = pricetag_pb2.PriceTag(
+            name="dollar",
+            cost=5.00,
+        )
+        self.assertIsNotNone(got)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/python/private/proto/py_proto_library.bzl b/python/private/proto/py_proto_library.bzl
index ef5f2ca..9885585 100644
--- a/python/private/proto/py_proto_library.bzl
+++ b/python/private/proto/py_proto_library.bzl
@@ -22,6 +22,9 @@
 _PyProtoInfo = provider(
     doc = "Encapsulates information needed by the Python proto rules.",
     fields = {
+        "imports": """
+            (depset[str]) The field forwarding PyInfo.imports coming from
+            the proto language runtime dependency.""",
         "runfiles_from_proto_deps": """
             (depset[File]) Files from the transitive closure implicit proto
             dependencies""",
@@ -95,6 +98,9 @@
 
     return [
         _PyProtoInfo(
+            imports = depset(
+                transitive = [dep[PyInfo].imports for dep in api_deps],
+            ),
             runfiles_from_proto_deps = runfiles_from_proto_deps,
             transitive_sources = transitive_sources,
         ),
@@ -142,6 +148,7 @@
         ),
         PyInfo(
             transitive_sources = default_outputs,
+            imports = depset(transitive = [info.imports for info in pyproto_infos]),
             # Proto always produces 2- and 3- compatible source files
             has_py2_only_sources = False,
             has_py3_only_sources = False,
diff --git a/tools/bazel_integration_test/bazel_integration_test.bzl b/tools/bazel_integration_test/bazel_integration_test.bzl
index 66e0cbd..c016551 100644
--- a/tools/bazel_integration_test/bazel_integration_test.bzl
+++ b/tools/bazel_integration_test/bazel_integration_test.bzl
@@ -84,18 +84,19 @@
     attrs = _ATTRS,
 )
 
-def bazel_integration_test(name, override_bazel_version = None, bzlmod = False, **kwargs):
+def bazel_integration_test(name, override_bazel_version = None, bzlmod = False, dirname = None, **kwargs):
     """Wrapper macro to set default srcs and run a py_test with config
 
     Args:
         name: name of the resulting py_test
         override_bazel_version: bazel version to use in test
         bzlmod: whether the test uses bzlmod
+        dirname: the directory name of the test. Defaults to value of `name` after trimming the `_example` suffix.
         **kwargs: additional attributes like timeout and visibility
     """
 
     # By default, we assume sources for "pip_example" are in examples/pip/**/*
-    dirname = name[:-len("_example")]
+    dirname = dirname or name[:-len("_example")]
     native.filegroup(
         name = "_%s_sources" % name,
         srcs = native.glob(