Copy py_proto_library from rules_python to protobuf

https://github.com/bazelbuild/rules_python/blob/main/python/private/proto/py_proto_library.bzl

Contributors:
d96214f [email protected]      Wed Nov 15 02:48:06 2023 -0800  fix: py_proto_library: transitive strip_import_prefix (#1558)
85e50d2 [email protected]       Tue Nov 14 06:04:59 2023 -0800  fix: py_proto_library: append to PYTHONPATH less (#1553)
bee35ef [email protected]  Wed Oct 11 20:59:34 2023 -0700  fix: allowing to import code generated from proto with strip_import_prefix (#1406)
1a333ce [email protected]        Tue Jun 20 19:36:39 2023 +0200  fix: plugin_output in py_proto_library rule (#1280)
6905e63 [email protected]   Sat Feb 11 14:02:33 2023 +0900  fix: make py_proto_library respect PyInfo imports (#1046)
0d3c4f7 [email protected]        Wed Jan 18 23:15:52 2023 +0000  Implement py_proto_library (#832)
PiperOrigin-RevId: 623401031

bazel/py_proto_library.bzl

@@ -0,0 +1,200 @@
+"""The implementation of the `py_proto_library` rule and its aspect."""
+
+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
+
+_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""",
+        "transitive_sources": """(depset[File]) The Python sources.""",
+    },
+)
+
+def _filter_provider(provider, *attrs):
+    return [dep[provider] for attr in attrs for dep in attr if provider in dep]
+
+def _py_proto_aspect_impl(target, ctx):
+    """Generates and compiles Python code for a proto_library.
+
+    The function runs protobuf compiler on the `proto_library` target generating
+    a .py file for each .proto file.
+
+    Args:
+      target: (Target) A target providing `ProtoInfo`. Usually this means a
+         `proto_library` target, but not always; you must expect to visit
+         non-`proto_library` targets, too.
+      ctx: (RuleContext) The rule context.
+
+    Returns:
+      ([_PyProtoInfo]) Providers collecting transitive information about
+      generated files.
+    """
+
+    _proto_library = ctx.rule.attr
+
+    # Check Proto file names
+    for proto in target[ProtoInfo].direct_sources:
+        if proto.is_source and "-" in proto.dirname:
+            fail("Cannot generate Python code for a .proto whose path contains '-' ({}).".format(
+                proto.path,
+            ))
+
+    proto_lang_toolchain_info = ctx.attr._aspect_proto_toolchain[ProtoLangToolchainInfo]
+    api_deps = [proto_lang_toolchain_info.runtime]
+
+    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(
+            actions = ctx.actions,
+            proto_info = proto_info,
+            extension = "_pb2.py",
+            name_mapper = lambda name: name.replace("-", "_").replace(".", "/"),
+        )
+
+        # Handles multiple repository and virtual import cases
+        if proto_root.startswith(ctx.bin_dir.path):
+            proto_root = proto_root[len(ctx.bin_dir.path) + 1:]
+
+        plugin_output = ctx.bin_dir.path + "/" + proto_root
+        proto_root = ctx.workspace_name + "/" + proto_root
+
+        proto_common.compile(
+            actions = ctx.actions,
+            proto_info = proto_info,
+            proto_lang_toolchain_info = proto_lang_toolchain_info,
+            generated_files = generated_sources,
+            plugin_output = plugin_output,
+        )
+
+    # Generated sources == Python sources
+    python_sources = generated_sources
+
+    deps = _filter_provider(_PyProtoInfo, getattr(_proto_library, "deps", []))
+    runfiles_from_proto_deps = depset(
+        transitive = [dep[DefaultInfo].default_runfiles.files for dep in api_deps] +
+                     [dep.runfiles_from_proto_deps for dep in deps],
+    )
+    transitive_sources = depset(
+        direct = python_sources,
+        transitive = [dep.transitive_sources for dep in deps],
+    )
+
+    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. But it's undesirable otherwise, because it
+                # will put the repo root at the top of the PYTHONPATH, ahead of
+                # directories added through `imports` attributes.
+                [proto_root] if "_virtual_imports" in proto_root else [],
+                transitive = [dep[PyInfo].imports for dep in api_deps] + [dep.imports for dep in deps],
+            ),
+            runfiles_from_proto_deps = runfiles_from_proto_deps,
+            transitive_sources = transitive_sources,
+        ),
+    ]
+
+_py_proto_aspect = aspect(
+    implementation = _py_proto_aspect_impl,
+    attrs = {
+        "_aspect_proto_toolchain": attr.label(
+            default = "//python:python_toolchain",
+        ),
+    },
+    attr_aspects = ["deps"],
+    required_providers = [ProtoInfo],
+    provides = [_PyProtoInfo],
+)
+
+def _py_proto_library_rule(ctx):
+    """Merges results of `py_proto_aspect` in `deps`.
+
+    Args:
+      ctx: (RuleContext) The rule context.
+    Returns:
+      ([PyInfo, DefaultInfo, OutputGroupInfo])
+    """
+    if not ctx.attr.deps:
+        fail("'deps' attribute mustn't be empty.")
+
+    pyproto_infos = _filter_provider(_PyProtoInfo, ctx.attr.deps)
+    default_outputs = depset(
+        transitive = [info.transitive_sources for info in pyproto_infos],
+    )
+
+    return [
+        DefaultInfo(
+            files = default_outputs,
+            default_runfiles = ctx.runfiles(transitive_files = depset(
+                transitive =
+                    [default_outputs] +
+                    [info.runfiles_from_proto_deps for info in pyproto_infos],
+            )),
+        ),
+        OutputGroupInfo(
+            default = depset(),
+        ),
+        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,
+        ),
+    ]
+
+py_proto_library = rule(
+    implementation = _py_proto_library_rule,
+    doc = """
+      Use `py_proto_library` to generate Python libraries from `.proto` files.
+
+      The convention is to name the `py_proto_library` rule `foo_py_pb2`,
+      when it is wrapping `proto_library` rule `foo_proto`.
+
+      `deps` must point to a `proto_library` rule.
+
+      Example:
+
+```starlark
+py_library(
+    name = "lib",
+    deps = [":foo_py_pb2"],
+)
+
+py_proto_library(
+    name = "foo_py_pb2",
+    deps = [":foo_proto"],
+)
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+)
+```""",
+    attrs = {
+        "deps": attr.label_list(
+            doc = """
+              The list of `proto_library` rules to generate Python libraries for.
+
+              Usually this is just the one target: the proto library of interest.
+              It can be any target providing `ProtoInfo`.""",
+            providers = [ProtoInfo],
+            aspects = [_py_proto_aspect],
+        ),
+    },
+    provides = [PyInfo],
+)

examples/BUILD.bazel

@@ -9,15 +9,17 @@ load("@protobuf//bazel:cc_proto_library.bzl", "cc_proto_library")
 load("@protobuf//bazel:java_lite_proto_library.bzl", "java_lite_proto_library")
 load("@protobuf//bazel:java_proto_library.bzl", "java_proto_library")
 load("@protobuf//bazel:proto_library.bzl", "proto_library")
+load("@protobuf//bazel:py_proto_library.bzl", "py_proto_library")
 load("@rules_cc//cc:defs.bzl", "cc_binary")
 load("@rules_pkg//pkg:mappings.bzl", "pkg_files", "strip_prefix")
+load("@rules_python//python:py_binary.bzl", "py_binary")
 
 # For each .proto file, a proto_library target should be defined. This target
 # is not bound to any particular language. Instead, it defines the dependency
 # graph of the .proto files (i.e., proto imports) and serves as the provider
 # of .proto source files to the protocol compiler.
 #
-# Remote repository "com_google_protobuf" must be defined to use this rule.
+# Remote repository "protobuf" must be defined to use this rule.
 proto_library(
     name = "addressbook_proto",
     srcs = ["addressbook.proto"],
@@ -116,11 +118,38 @@ java_binary(
     deps = [":addressbook_java_lite_proto"],
 )
 
+# Python
+
+py_proto_library(
+    name = "addressbook_py_pb2",
+    visibility = ["//visibility:public"],
+    deps = [":addressbook_proto"],
+)
+
+py_binary(
+    name = "add_person",
+    srcs = ["add_person.py"],
+    python_version = "PY3",
+    deps = [
+        ":addressbook_py_pb2",
+    ],
+)
+
+py_binary(
+    name = "list_people",
+    srcs = ["list_people.py"],
+    python_version = "PY3",
+    deps = [
+        ":addressbook_py_pb2",
+    ],
+)
+
 build_test(
     name = "test",
     targets = [
         ":add_person_cpp",
         ":add_person_java",
+        ":add_person",  # Python
     ],
 )
 

examples/MODULE.bazel

@@ -10,3 +10,4 @@ bazel_dep(name = "bazel_skylib", version = "1.0.3")
 bazel_dep(name = "rules_cc", version = "0.0.1")
 bazel_dep(name = "rules_java", version = "7.3.0")
 bazel_dep(name = "rules_pkg", version = "0.7.0")
+bazel_dep(name = "rules_python", version = "0.25.0")

python/build_targets.bzl

@@ -9,6 +9,7 @@
 load("@rules_pkg//pkg:mappings.bzl", "pkg_files", "strip_prefix")
 load("@rules_python//python:defs.bzl", "py_library")
 load("//:protobuf.bzl", "internal_py_proto_library")
+load("//bazel/toolchains:proto_lang_toolchain.bzl", "proto_lang_toolchain")
 load("//build_defs:arch_tests.bzl", "aarch64_test", "x86_64_test")
 load("//build_defs:cpp_opts.bzl", "COPTS")
 load("//conformance:defs.bzl", "conformance_test")
@@ -510,3 +511,13 @@ def build_targets(name):
         strip_prefix = strip_prefix.from_root(""),
         visibility = ["//pkg:__pkg__"],
     )
+
+    proto_lang_toolchain(
+        name = "python_toolchain",
+        command_line = "--python_out=%s",
+        progress_message = "Generating Python proto_library %{label}",
+        runtime = ":protobuf_python",
+        # NOTE: This isn't *actually* public. It's an implicit dependency of py_proto_library,
+        # so must be public so user usages of the rule can reference it.
+        visibility = ["//visibility:public"],
+    )