WIP: start proto/gRPC support

.bazelrc

@@ -10,8 +10,10 @@ common --define=SOME_VAR=SOME_VALUE
 common --@pypi//venv=default
 
 common --incompatible_enable_cc_toolchain_resolution
+common --incompatible_enable_proto_toolchain_resolution
 common --@toolchains_llvm_bootstrapped//config:experimental_stub_libgcc_s
 common --@rules_cc//cc/toolchains/args/archiver_flags:use_libtool_on_macos=false
+common --@protobuf//bazel/toolchains:prefer_prebuilt_protoc
 
 # TODO(bzlmod): Don't break proto
 common --per_file_copt=external/.*protobuf.*@--PROTOBUF_WAS_NOT_SUPPOSED_TO_BE_BUILT

MODULE.bazel

@@ -18,6 +18,9 @@ bazel_dep(name = "rules_cc", version = "0.2.16")
 bazel_dep(name = "rules_pkg", version = "1.1.0")
 bazel_dep(name = "tar.bzl", version = "0.5.5")
 
+# NB: LOWER BOUND on earliest BCR release of protobuf module, to avoid upgrading the root module by accident
+bazel_dep(name = "protobuf", version = "3.19.6")
+
 bazel_lib_toolchains = use_extension("@tar.bzl//tar:extensions.bzl", "toolchains")
 use_repo(bazel_lib_toolchains, "bsd_tar_toolchains")
 
@@ -103,6 +106,17 @@ bazel_dep(name = "xz", version = "5.4.5.bcr.7", dev_dependency = IS_RELEASE)
 bazel_dep(name = "zstd", version = "1.5.7", dev_dependency = IS_RELEASE)
 bazel_dep(name = "rules_multitool", version = "1.9.0", dev_dependency = IS_RELEASE)
 
+# Support examples/protobuf
+single_version_override(
+    module_name = "protobuf",
+    version = "33.4",
+)
+
+register_toolchains(
+    "//examples/protobuf/toolchains:all",
+    dev_dependency = True,
+)
+
 multitool = use_extension("@rules_multitool//multitool:extension.bzl", "multitool", dev_dependency = IS_RELEASE)
 multitool.hub(lockfile = "//tools:tools.lock.json")
 use_repo(multitool, "multitool")

examples/protobuf/BUILD.bazel

@@ -0,0 +1,13 @@
+load("@aspect_rules_py//py:defs.bzl", "py_binary")
+load("@protobuf//bazel:proto_library.bzl", "proto_library")
+
+proto_library(
+    name = "foo_proto",
+    srcs = ["foo.proto"],
+)
+
+py_binary(
+    name = "main",
+    srcs = ["__main__.py"],
+    deps = [":foo_proto"],
+)

examples/protobuf/__main__.py

@@ -0,0 +1,6 @@
+from examples.protobuf import foo_pb2
+
+foo = foo_pb2.Foo()
+foo.name = "Hello, World!"
+
+print(foo)

examples/protobuf/foo.proto

@@ -0,0 +1,7 @@
+syntax = "proto3";
+
+package foo;
+
+message Foo {
+  string name = 1;
+}

examples/protobuf/foo_pb2.pyi

@@ -0,0 +1,11 @@
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Optional as _Optional
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class Foo(_message.Message):
+    __slots__ = ("name",)
+    NAME_FIELD_NUMBER: _ClassVar[int]
+    name: str
+    def __init__(self, name: _Optional[str] = ...) -> None: ...

examples/protobuf/toolchains/BUILD.bazel

@@ -0,0 +1,26 @@
+load("@aspect_rules_py//py:proto.bzl", "py_proto_toolchain")
+
+# FIXME(arrdem): does rules_py uv-created package support any way to invoke the console script of a package?
+# py_console_script_binary(
+#     name = "protoc-gen-grpc",
+#     pkg = "@pypi//grpcio_tools",
+#     script = "protoc",
+# )
+
+py_proto_toolchain(
+    name = "builtin_protoc_plugin",
+    plugin_bin = None,  # use the python emit embedded in protoc
+    plugin_name = "python",
+    plugin_options = [],
+    runtime = "@pypi//protobuf",
+)
+
+# TODO: support gRPC as well, which is NOT built-in
+# py_proto_toolchain(
+#     name = "grpcio_protoc_plugin",
+#     plugin_bin = ":protoc-gen-grpc",
+#     plugin_name = "python",
+#     plugin_options = [
+#     ],
+#     # runtime = "//:node_modules/@bufbuild/protobuf",
+# )

py/BUILD.bazel

@@ -52,3 +52,13 @@ bzl_library(
         "@bazel_tools//tools/build_defs/repo:http.bzl",
     ],
 )
+
+bzl_library(
+    name = "proto",
+    srcs = ["proto.bzl"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//py/private:proto",
+        "@protobuf//bazel/toolchains:proto_lang_toolchain_bzl",
+    ],
+)

py/private/proto.bzl

@@ -0,0 +1,98 @@
+"""An aspect that generates Python code from .proto files.
+
+The aspect converts a ProtoInfo provider into a PyInfo provider so that proto_library may be a dep to python rules.
+"""
+
+load("@protobuf//bazel/common:proto_common.bzl", "proto_common")
+load("@protobuf//bazel/common:proto_info.bzl", "ProtoInfo")
+load("@rules_python//python:defs.bzl", "PyInfo")
+
+LANG_PROTO_TOOLCHAIN = Label("//py/private/toolchain:protoc_plugin_toolchain_type")
+
+def _py_proto_aspect_impl(target, ctx):
+    proto_info = target[ProtoInfo]
+    proto_lang_toolchain_info = ctx.toolchains[LANG_PROTO_TOOLCHAIN].proto
+    proto_deps = [d for d in ctx.rule.attr.deps if PyInfo in d]
+    python_naming = lambda name: name.replace("-", "_").replace(".", "/")
+    py_outputs = proto_common.declare_generated_files(
+        actions = ctx.actions,
+        proto_info = proto_info,
+        extension = "_pb2.py",
+        name_mapper = python_naming,
+    )
+
+    generated_stubs = []
+    if True:  # len([o for o in proto_lang_toolchain_info.command_line.split() if o.startswith("--pyi_out=")]) > 0:
+        generated_stubs = proto_common.declare_generated_files(
+            actions = ctx.actions,
+            proto_info = proto_info,
+            extension = "_pb2.pyi",
+            name_mapper = python_naming,
+        )
+
+    # Determine root folder, mapping output paths to inputs, i.e. bazel-bin/arch/bin/foo to foo
+    proto_root = proto_info.proto_source_root
+    if proto_root.startswith(ctx.bin_dir.path):
+        proto_root = proto_root[len(ctx.bin_dir.path) + 1:]
+
+    # FIXME: this is plugin-specific and fishy.
+    # the user should specify their pyi_out preference in the lang_proto_toolchain construction
+    additional_args = ctx.actions.args()
+    additional_args.add(py_outputs[0].root.path, format = "--pyi_out=%s")
+
+    # It's possible for proto_library to have only deps but no srcs
+    if proto_info.direct_sources:
+        proto_common.compile(
+            actions = ctx.actions,
+            proto_info = proto_info,
+            proto_lang_toolchain_info = proto_lang_toolchain_info,
+            generated_files = py_outputs + generated_stubs,
+            plugin_output = py_outputs[0].root.path,
+            additional_args = additional_args,
+        )
+
+    # Import path within the runfiles tree
+    if proto_root.startswith("external/"):
+        import_path = proto_root[len("external") + 1:]
+    else:
+        import_path = ctx.workspace_name + "/" + proto_root
+    return [
+        DefaultInfo(files = depset(generated_stubs)),
+        PyInfo(
+            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.
+                [import_path] if "_virtual_imports" in import_path else [],
+                transitive = [dep.imports for dep in proto_deps] + (
+                    [proto_lang_toolchain_info.runtime[PyInfo].imports] if proto_lang_toolchain_info.runtime else []
+                ),
+            ),
+            # direct_pyi_files = depset(direct = direct_pyi_files),
+            # transitive_pyi_files = transitive_pyi_files,
+            transitive_sources = depset(
+                py_outputs + generated_stubs,
+                transitive = [dep.transitive_sources for dep in proto_deps] + (
+                    [proto_lang_toolchain_info.runtime[PyInfo].transitive_sources] if proto_lang_toolchain_info.runtime else []
+                ),
+            ),
+            # Proto always produces 2- and 3- compatible source files
+            has_py2_only_sources = False,
+            has_py3_only_sources = False,
+            uses_shared_libraries = False,
+        ),
+    ]
+
+py_proto_aspect = aspect(
+    implementation = _py_proto_aspect_impl,
+    # Traverse the "deps" graph edges starting from the target
+    attr_aspects = ["deps"],
+    # Only visit nodes that produce a ProtoInfo provider
+    required_providers = [ProtoInfo],
+    # Be a valid dependency of a py_library rule
+    provides = [PyInfo],
+    toolchains = [LANG_PROTO_TOOLCHAIN],
+)

py/private/py_library.bzl

@@ -9,6 +9,7 @@ load("@bazel_skylib//lib:paths.bzl", "paths")
 load("@rules_cc//cc/common:cc_info.bzl", "CcInfo")
 load("@rules_python//python:defs.bzl", "PyInfo")
 load("//py/private:providers.bzl", "PyVirtualInfo")
+load(":proto.bzl", "py_proto_aspect")
 
 def _make_instrumented_files_info(ctx, extra_source_attributes = [], extra_dependency_attributes = []):
     return coverage_common.instrumented_files_info(
@@ -206,6 +207,7 @@ _attrs = dict({
     "deps": attr.label_list(
         doc = "Targets that produce Python code, commonly `py_library` rules.",
         providers = [[PyInfo], [PyVirtualInfo], [CcInfo]],
+        aspects = [py_proto_aspect],
     ),
     "data": attr.label_list(
         doc = """Runtime dependencies of the program.

py/private/toolchain/BUILD.bazel

@@ -6,6 +6,11 @@ exports_files(
     visibility = ["//visibility:public"],
 )
 
+toolchain_type(
+    name = "protoc_plugin_toolchain_type",
+    visibility = ["//visibility:public"],
+)
+
 toolchain_type(
     name = "unpack_toolchain_type",
     visibility = ["//visibility:public"],

py/proto.bzl

@@ -0,0 +1,101 @@
+"""**EXPERIMENTAL**: Protobuf and gRPC support for Python.
+
+This API is subject to breaking changes outside our usual semver policy.
+In a future release of rules_py this should become stable.
+
+### Typical setup
+
+1. Choose any code generator plugin for protoc.
+   For example, https://pypi.org/project/grpcio-tools/ which provides `protoc` and
+   Python gRPC code generation.
+2. Declare a binary target that runs the generator, for example:
+
+```starlark
+load("@rules_python//python/entry_points:py_console_script_binary.bzl", "py_console_script_binary")
+
+py_console_script_binary(
+    name = "protoc-gen-grpc",
+    pkg = "@pypi//grpcio_tools",
+    script = "protoc",
+)
+```
+3. Instead, it's also possible to use the Python plugin which is built-in to protoc.
+3. Define a `py_proto_toolchain` that specifies the plugin. See the rule documentation below.
+4. Update `MODULE.bazel` to register it, typically with
+   `register_toolchains("//tools/toolchains:all")`.
+
+### Usage
+
+Write `proto_library` targets as usual, or have Gazelle generate them.
+Then reference them anywhere a `py_library` could appear.
+Note this attribute is not supported by rules_python, meaning your BUILD files will be specific to rules_py.
+
+For example:
+
+```starlark
+load("@aspect_rules_py//py:defs.bzl", "py_library")
+load("@protobuf//bazel:proto_library.bzl", "proto_library")
+
+proto_library(
+    name = "eliza_proto",
+    srcs = ["eliza.proto"],
+)
+
+py_library(
+    name = "proto",
+    deps = [":eliza_proto"],
+)
+```
+"""
+
+load("@protobuf//bazel/toolchains:proto_lang_toolchain.bzl", "proto_lang_toolchain")
+load("//py/private:proto.bzl", "LANG_PROTO_TOOLCHAIN")
+
+def py_proto_toolchain(name, plugin_name, plugin_options, plugin_bin, runtime = None, **kwargs):
+    """Define a proto_lang_toolchain that uses the plugin.
+
+    Example:
+
+    ```starlark
+    py_proto_toolchain(
+        name = "gen_es_protoc_plugin",
+        plugin_bin = ":protoc-gen-grpc",
+        plugin_name = "python",
+        plugin_options = [
+        ],
+        runtime = "@pypi//:protobuf",
+    )
+    ```
+
+    Args:
+        name: The name of the toolchain. A target named [name]_toolchain is also created, which is the one to be used in register_toolchains.
+        plugin_name: The `NAME` of the plugin program, used in command-line flags to protoc, as follows:
+
+            > `protoc --plugin=protoc-gen-NAME=path/to/mybinary --NAME_out=OUT_DIR`
+
+            See https://protobuf.dev/reference/cpp/api-docs/google.protobuf.compiler.plugin
+
+        plugin_options: (List of strings) Command line flags used to invoke the plugin,
+            based on documentation for the generator.
+
+        plugin_bin: The plugin to use. This should be the label of a binary target that you declared in step 2 above.
+            If `None`, the Python plugin which is built-in to protoc will be used.
+
+        runtime: Optional runtime dependency imported by generated code.
+
+        **kwargs: Additional arguments to pass to the [proto_lang_toolchain](https://bazel.build/reference/be/protocol-buffer#proto_lang_toolchain) rule.
+    """
+    command_line_flags = ["--{}_opt=%s".format(plugin_name) % o for o in plugin_options]
+    command_line_flags.append("--{}_out=$(OUT)".format(plugin_name))
+    attrs = dict(kwargs)
+    if runtime != None:
+        attrs["runtime"] = runtime
+
+    proto_lang_toolchain(
+        name = name,
+        command_line = " ".join(command_line_flags),
+        plugin_format_flag = "--plugin=protoc-gen-{}=%s".format(plugin_name),
+        toolchain_type = LANG_PROTO_TOOLCHAIN,
+        plugin = plugin_bin,
+        **attrs
+    )

pyproject.toml

@@ -0,0 +1,4 @@
+[tool.pyright]
+# Bazel-generated modules may provide `.pyi` in source, while `.py` appears only
+# in runfiles/venv at runtime. Suppress missing-source diagnostics for this case.
+reportMissingModuleSource = "none"

requirements.in

@@ -11,3 +11,5 @@ ftfy==6.2.0
 neptune==1.10.2
 six
 bazel-runfiles
+grpcio-tools
+protobuf==6.33.4