diff --git a/src/python/pants/backend/codegen/protobuf/buf/BUILD b/src/python/pants/backend/codegen/protobuf/buf/BUILD new file mode 100644 index 00000000000..4288a120636 --- /dev/null +++ b/src/python/pants/backend/codegen/protobuf/buf/BUILD @@ -0,0 +1,6 @@ +# Copyright 2026 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +python_sources() + +python_tests(name="tests") diff --git a/src/python/pants/backend/codegen/protobuf/buf/__init__.py b/src/python/pants/backend/codegen/protobuf/buf/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/src/python/pants/backend/codegen/protobuf/buf/config.py b/src/python/pants/backend/codegen/protobuf/buf/config.py new file mode 100644 index 00000000000..dc3e4fc797e --- /dev/null +++ b/src/python/pants/backend/codegen/protobuf/buf/config.py @@ -0,0 +1,454 @@ +# Copyright 2026 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +"""Shared, language-agnostic helpers for working with `buf.yaml` and `buf.gen.yaml`. + +These primitives are used by the per-language buf integrations so each language +only owns its own suffix conventions and module-name math, not yaml parsing +or plugin matching. +""" + +from __future__ import annotations + +import os +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + +import yaml + +from pants.backend.codegen.protobuf.buf.fields import BufGenTemplateField +from pants.backend.codegen.protobuf.buf.subsystem import BufSubsystem +from pants.core.util_rules.config_files import ConfigFilesRequest, find_config_file +from pants.engine.intrinsics import get_digest_contents +from pants.engine.rules import concurrently +from pants.engine.target import Target + +# ---- Errors ---------------------------------------------------------------- + + +class UnpinnedBufPluginError(Exception): + """Raised when a `remote:` plugin entry is missing a version+revision pin and + isn't in the user's `DEFAULT_PLUGIN_PINS` registry.""" + + +class MissingBufLockError(Exception): + """Raised when `buf.yaml` declares `deps:` but no sibling `buf.lock` exists.""" + + +# ---- Default plugin-pin registry ------------------------------------------- + + +# Default `(version, revision)` pin Pants will fill in for known `remote:` +# plugins that the user wrote without a pin or revision. The synthesized +# `buf.gen.yaml` entry has the form: +# +# - remote: : +# revision: +# out: ... +# +# These values are the latest as of writing, fetched from the BSR's +# `PluginCurationService.GetLatestCuratedPlugin` endpoint. To refresh: +# +# curl -X POST \ +# https://buf.build/buf.alpha.registry.v1alpha1.PluginCurationService/GetLatestCuratedPlugin \ +# -H 'Content-Type: application/json' -H 'Connect-Protocol-Version: 1' \ +# -d '{"owner":"","name":""}' +# +# The browseable equivalent is https://buf.build// — each +# plugin's "Versions" tab lists every `(version, revision)` pair available. +# Bumping a pin is a one-line change here; cache invalidation is automatic +# via the synthesized `buf.gen.yaml`'s content hash. +DEFAULT_PLUGIN_PINS: Mapping[str, tuple[str, int]] = { + "buf.build/protocolbuffers/python": ("v34.1", 1), + "buf.build/protocolbuffers/pyi": ("v34.1", 1), + "buf.build/connectrpc/python": ("v0.10.0", 1), + "buf.build/grpc/python": ("v1.80.0", 1), +} + + +# ---- buf.yaml parsers ------------------------------------------------------ + + +def parse_buf_yaml_module_paths(content: bytes) -> tuple[str, ...]: + """Module paths declared in a v2 `buf.yaml`, relative to the file. + + Returns `()` for v1 `buf.yaml` (no `modules:` block); callers should treat the + file's own directory as the implicit module root in that case. + """ + try: + data = yaml.safe_load(content) + except yaml.YAMLError: + return () + if not isinstance(data, dict): + return () + modules = data.get("modules") + if not isinstance(modules, list): + return () + paths: list[str] = [] + for entry in modules: + if isinstance(entry, dict): + p = entry.get("path") + if isinstance(p, str) and p: + paths.append(p) + return tuple(paths) + + +def parse_buf_yaml_deps(content: bytes) -> tuple[str, ...]: + """BSR module IDs declared in a `buf.yaml`'s `deps:` list.""" + try: + data = yaml.safe_load(content) + except yaml.YAMLError: + return () + if not isinstance(data, dict): + return () + deps = data.get("deps") + if not isinstance(deps, list): + return () + return tuple(d for d in deps if isinstance(d, str) and d) + + +def resolve_buf_module_root( + proto_path: str, buf_yaml_dir: str, module_paths: tuple[str, ...] +) -> str: + """Resolve the buf module root for a proto file. + + `buf_yaml_dir` is the directory holding `buf.yaml`. `module_paths` are the + entries from its `modules:` list (relative to that directory). The returned + root is the longest configured module path that contains the proto, normalized + to the repo root. + """ + if not module_paths: + return buf_yaml_dir + candidates = sorted( + (os.path.normpath(os.path.join(buf_yaml_dir, p)) for p in module_paths), + key=len, + reverse=True, + ) + for root in candidates: + if root == "." or root == "": + return "" + if proto_path == root or proto_path.startswith(root + os.sep): + return root + return candidates[0] + + +# ---- buf.gen.yaml parsers -------------------------------------------------- + + +def _plugin_identifier(plugin: dict) -> str | None: + """Return the registry key (`":"`) for a `buf.gen.yaml` plugin entry. + + `kind` is one of `protoc_builtin`, `local`, or `remote` — matching the field + name in the entry. `ident` strips any `:vX.Y` version pin from `remote:` so + callers can match against the registry without knowing the version. Returns + `None` if the entry declares none of these fields. + """ + for key in ("protoc_builtin", "local", "remote"): + val = plugin.get(key) + if val is None: + continue + ident = " ".join(str(x) for x in val) if isinstance(val, list) else str(val) + if key == "remote": + base, _, _ = ident.partition(":") + return f"remote:{base}" + return f"{key}:{ident}" + return None + + +def parse_plugin_outs(content: bytes, suffixes: Mapping[str, str]) -> dict[str, str]: + """Walk `buf.gen.yaml` plugins and return `suffix -> out:` for matching entries. + + `suffixes` is a `: -> suffix` dict supplied by the calling + language backend, where `` is `remote`, `protoc_builtin`, or `local` + and `suffix` is the language's module/file-naming suffix. The first matching + plugin per suffix wins. + """ + try: + data = yaml.safe_load(content) + except yaml.YAMLError: + return {} + if not isinstance(data, dict): + return {} + plugins = data.get("plugins") + if not isinstance(plugins, list): + return {} + + result: dict[str, str] = {} + for plugin in plugins: + if not isinstance(plugin, dict): + continue + out = plugin.get("out") + if not isinstance(out, str) or not out: + continue + key = _plugin_identifier(plugin) + if key is None: + continue + suffix = suffixes.get(key) + if suffix is not None and suffix not in result: + result[suffix] = out + return result + + +def suffix_plugin_includes_imports( + content: bytes, suffix: str, suffixes: Mapping[str, str] +) -> bool: + """True if the `buf.gen.yaml` plugin emitting `suffix` has `include_imports: + true` set — meaning buf will materialize generated artifacts for + transitively-imported BSR-dep protos into the digest. + + `suffixes` is the same `: -> suffix` mapping passed to + `parse_plugin_outs`; the first plugin entry whose registry key maps to + `suffix` wins. + """ + try: + data = yaml.safe_load(content) + except yaml.YAMLError: + return False + if not isinstance(data, dict): + return False + plugins = data.get("plugins") + if not isinstance(plugins, list): + return False + for plugin in plugins: + if not isinstance(plugin, dict): + continue + out = plugin.get("out") + if not isinstance(out, str) or not out: + continue + key = _plugin_identifier(plugin) + if key is None: + continue + if suffixes.get(key) != suffix: + continue + return plugin.get("include_imports") is True + return False + + +# ---- buf.gen.yaml pin synthesis -------------------------------------------- + + +def _split_remote_ident(ident: str) -> tuple[str, str | None]: + """Split a `remote:` value into `(base_id, version_or_None)`. + + `buf.build/foo/bar:v1.2` → (`buf.build/foo/bar`, `v1.2`). + `buf.build/foo/bar` → (`buf.build/foo/bar`, None). + """ + base, sep, suffix = ident.partition(":") + return (base, suffix) if sep else (base, None) + + +def _parse_pin_string(pin: str) -> tuple[str, int] | None: + """Parse `"vX.Y:N"` (Pants-internal pin format) into `(version, revision)`. + Returns `None` if the format is invalid.""" + parts = pin.split(":") + if len(parts) != 2 or not parts[0]: + return None + try: + return parts[0], int(parts[1]) + except ValueError: + return None + + +def synthesize_pinned_buf_gen_yaml( + content: bytes, + source_path: str, + *, + extra_pins: Mapping[str, str] | None = None, +) -> bytes: + """Return the user's `buf.gen.yaml` with `remote:` plugin pins resolved. + + For each `remote:` entry, the buf-recognized pinned form is: + + - remote: : + revision: + out: ... + + Pants requires both fields. If a plugin entry is missing either, it must be + in `DEFAULT_PLUGIN_PINS` (or the user's `extra_pins`) for Pants to fill in + defaults. `extra_pins` values are `"vX.Y:N"` strings, parsed here. + `protoc_builtin:` and `local:` plugins are not subject to pin enforcement. + """ + try: + data = yaml.safe_load(content) + except yaml.YAMLError: + return content + if not isinstance(data, dict): + return content + plugins = data.get("plugins") + if not isinstance(plugins, list): + return content + + parsed_extra: dict[str, tuple[str, int]] = {} + for k, v in (extra_pins or {}).items(): + parsed = _parse_pin_string(v) + if parsed is not None: + parsed_extra[k] = parsed + pins: Mapping[str, tuple[str, int]] = {**DEFAULT_PLUGIN_PINS, **parsed_extra} + + unresolvable: list[str] = [] + rewrote = False + for plugin in plugins: + if not isinstance(plugin, dict): + continue + val = plugin.get("remote") + if val is None: + continue + ident = " ".join(str(x) for x in val) if isinstance(val, list) else str(val) + base, version = _split_remote_ident(ident) + revision = plugin.get("revision") + if version is not None and isinstance(revision, int): + continue # already fully pinned + default = pins.get(base) + if default is None: + unresolvable.append(ident) + continue + default_version, default_revision = default + plugin["remote"] = f"{base}:{default_version}" + plugin["revision"] = default_revision + rewrote = True + + if unresolvable: + bullets = "\n".join(f" - remote: {ident}" for ident in unresolvable) + known = ", ".join(sorted(DEFAULT_PLUGIN_PINS)) or "(none)" + raise UnpinnedBufPluginError( + f"`{source_path}` has `remote:` plugin entries that are missing a " + f"version or `revision:`:\n{bullets}\n\n" + "Pin both fields explicitly:\n\n" + " - remote: buf.build/owner/plugin:vX.Y\n" + " revision: N\n" + " out: ...\n\n" + "Alternatively, for a plugin in Pants's built-in registry, leave both " + "fields unset and Pants will fill in defaults. To extend the registry " + "for your own plugins, use `[buf].extra_plugin_pins`.\n\n" + f"Built-in registry: {known}." + ) + + if not rewrote: + return content + return yaml.safe_dump(data, sort_keys=False).encode("utf-8") + + +def check_pinned_remote_plugins( + content: bytes, + source_path: str, + *, + extra_pins: Mapping[str, str] | None = None, +) -> None: + """Raise if `remote:` plugin entries can't be resolved to a full pin, without + returning the synthesized content.""" + synthesize_pinned_buf_gen_yaml(content, source_path, extra_pins=extra_pins) + + +# ---- Per-target template request resolvers -------------------------------- + + +def gen_template_request_from_fields( + *, + spec_path: str, + address_str: str, + override: str | None, + buf: BufSubsystem, +) -> ConfigFilesRequest: + """Resolve the `buf.gen.yaml` request from already-extracted field values. + + Precedence: per-target `buf_gen_template` (`override`) → `[buf].gen_template` + subsystem option → `[buf].gen_template_discovery`. + """ + if override is None: + return buf.gen_template_request + path = os.path.normpath(os.path.join(spec_path, override)) + return ConfigFilesRequest( + specified=path, + specified_option_name=f"`{BufGenTemplateField.alias}` field on {address_str}", + discovery=False, + check_existence=(path,), + ) + + +def gen_template_request_for_target(tgt: Target, buf: BufSubsystem) -> ConfigFilesRequest: + """Convenience wrapper around `gen_template_request_from_fields` for a Target.""" + return gen_template_request_from_fields( + spec_path=tgt.address.spec_path, + address_str=str(tgt.address), + override=tgt.get(BufGenTemplateField).value, + buf=buf, + ) + + +def resolved_template_path(tgt: Target, buf: BufSubsystem) -> str | None: + """Path to pass to `buf generate --template`, or None to rely on discovery.""" + override = tgt.get(BufGenTemplateField).value + if override is not None: + return os.path.normpath(os.path.join(tgt.address.spec_path, override)) + return buf.gen_template + + +# ---- Async fetchers + result types ---------------------------------------- + + +@dataclass(frozen=True) +class BufLayout: + """Module layout derived from `buf.yaml`.""" + + buf_yaml_dir: str + module_paths: tuple[str, ...] + deps: tuple[str, ...] # BSR module ids (e.g. `buf.build/bufbuild/protovalidate`) + + def root_for_proto(self, proto_path: str) -> str: + return resolve_buf_module_root(proto_path, self.buf_yaml_dir, self.module_paths) + + +async def fetch_buf_layout(buf: BufSubsystem) -> BufLayout: + """Read `buf.yaml` and return the parsed module layout. Empty if not found. + + `config_request` may also surface `buf.lock` (it's listed in `check_existence` + so codegen invalidates on lock changes), so we filter to `buf.yaml` here. + """ + files = await find_config_file(buf.config_request) + yaml_paths = [p for p in files.snapshot.files if os.path.basename(p) == "buf.yaml"] + if not yaml_paths: + return BufLayout("", (), ()) + path = yaml_paths[0] + contents = await get_digest_contents(files.snapshot.digest) + content = next((dc.content for dc in contents if dc.path == path), b"") + return BufLayout( + os.path.dirname(path), + parse_buf_yaml_module_paths(content), + parse_buf_yaml_deps(content), + ) + + +@dataclass(frozen=True) +class BufGenContent: + """Per-target `buf.gen.yaml` resolution. + + `template_path` is `None` when no template was found (callers should fall back + to source-root path arithmetic). When set, `content` is the raw yaml. + """ + + target: Target + template_path: str | None + content: bytes + + +async def fetch_buf_gen_contents( + targets: Sequence[Target], buf: BufSubsystem +) -> tuple[BufGenContent, ...]: + """Resolve and read each target's effective `buf.gen.yaml`.""" + if not targets: + return () + template_files_per_target = await concurrently( + find_config_file(gen_template_request_for_target(t, buf)) for t in targets + ) + contents_per_target = await concurrently( + get_digest_contents(tf.snapshot.digest) for tf in template_files_per_target + ) + out: list[BufGenContent] = [] + for tgt, files, dcs in zip(targets, template_files_per_target, contents_per_target): + if not files.snapshot.files: + out.append(BufGenContent(tgt, None, b"")) + continue + path = files.snapshot.files[0] + content = next((dc.content for dc in dcs if dc.path == path), b"") + out.append(BufGenContent(tgt, path, content)) + return tuple(out) diff --git a/src/python/pants/backend/codegen/protobuf/buf/config_test.py b/src/python/pants/backend/codegen/protobuf/buf/config_test.py new file mode 100644 index 00000000000..51a52eae150 --- /dev/null +++ b/src/python/pants/backend/codegen/protobuf/buf/config_test.py @@ -0,0 +1,289 @@ +# Copyright 2026 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +from __future__ import annotations + +from textwrap import dedent + +import pytest +import yaml + +from pants.backend.codegen.protobuf.buf.config import ( + UnpinnedBufPluginError, + parse_buf_yaml_deps, + parse_plugin_outs, + suffix_plugin_includes_imports, + synthesize_pinned_buf_gen_yaml, +) + + +def _plugins(content: bytes) -> list[dict]: + parsed = yaml.safe_load(content) + assert isinstance(parsed, dict) + plugins = parsed["plugins"] + assert isinstance(plugins, list) + return plugins + + +def test_synthesize_keeps_already_pinned_entry_unchanged() -> None: + content = dedent( + """\ + version: v2 + plugins: + - remote: buf.build/protocolbuffers/python:v34.1 + revision: 1 + out: gen + """ + ).encode("utf-8") + out = synthesize_pinned_buf_gen_yaml(content, "buf.gen.yaml") + # Already fully pinned → returned unchanged byte-for-byte. + assert out == content + + +def test_synthesize_fills_in_pin_for_known_unpinned_plugin() -> None: + content = dedent( + """\ + version: v2 + plugins: + - remote: buf.build/protocolbuffers/python + out: gen + """ + ).encode("utf-8") + out = synthesize_pinned_buf_gen_yaml(content, "buf.gen.yaml") + [plugin] = _plugins(out) + # Registry default for protocolbuffers/python kicks in. + assert plugin["remote"].startswith("buf.build/protocolbuffers/python:v") + assert isinstance(plugin["revision"], int) and plugin["revision"] >= 1 + + +def test_synthesize_overrides_partial_pin_with_registry_default() -> None: + """A version-only pin (no `revision:`) gets overridden with the registry's full + pin so the entry is unambiguous.""" + content = dedent( + """\ + version: v2 + plugins: + - remote: buf.build/protocolbuffers/python:v33.0 + out: gen + """ + ).encode("utf-8") + out = synthesize_pinned_buf_gen_yaml(content, "buf.gen.yaml") + [plugin] = _plugins(out) + assert "revision" in plugin + # Registry default version replaces the user's version-only pin. + assert plugin["remote"] != "buf.build/protocolbuffers/python:v33.0" + + +def test_synthesize_raises_for_unknown_unpinned_plugin() -> None: + content = dedent( + """\ + version: v2 + plugins: + - remote: example.com/some/custom-plugin + out: gen + """ + ).encode("utf-8") + with pytest.raises(UnpinnedBufPluginError) as excinfo: + synthesize_pinned_buf_gen_yaml(content, "buf.gen.yaml") + assert "example.com/some/custom-plugin" in str(excinfo.value) + + +def test_synthesize_accepts_user_provided_extra_pins() -> None: + content = dedent( + """\ + version: v2 + plugins: + - remote: example.com/some/custom-plugin + out: gen + """ + ).encode("utf-8") + out = synthesize_pinned_buf_gen_yaml( + content, + "buf.gen.yaml", + extra_pins={"example.com/some/custom-plugin": "v2.0:3"}, + ) + [plugin] = _plugins(out) + assert plugin["remote"] == "example.com/some/custom-plugin:v2.0" + assert plugin["revision"] == 3 + + +def test_synthesize_extra_pins_override_registry_default() -> None: + """A user `extra_pins` entry for a plugin already in the registry overrides + the registry default.""" + content = dedent( + """\ + version: v2 + plugins: + - remote: buf.build/protocolbuffers/python + out: gen + """ + ).encode("utf-8") + out = synthesize_pinned_buf_gen_yaml( + content, + "buf.gen.yaml", + extra_pins={"buf.build/protocolbuffers/python": "v99.0:7"}, + ) + [plugin] = _plugins(out) + assert plugin["remote"] == "buf.build/protocolbuffers/python:v99.0" + assert plugin["revision"] == 7 + + +def test_synthesize_ignores_malformed_extra_pin_string() -> None: + """Malformed `extra_pins` values (missing `:revN`) are silently dropped, so the + plugin is treated as unknown-and-unpinned and raises.""" + content = dedent( + """\ + version: v2 + plugins: + - remote: example.com/some/custom-plugin + out: gen + """ + ).encode("utf-8") + with pytest.raises(UnpinnedBufPluginError): + synthesize_pinned_buf_gen_yaml( + content, + "buf.gen.yaml", + extra_pins={"example.com/some/custom-plugin": "no-colon"}, + ) + + +def test_synthesize_ignores_protoc_builtin_and_local() -> None: + content = dedent( + """\ + version: v2 + plugins: + - protoc_builtin: python + out: gen + - local: protoc-gen-foo + out: gen + """ + ).encode("utf-8") + # No `remote:` entries → no synthesis, no error. + assert synthesize_pinned_buf_gen_yaml(content, "buf.gen.yaml") == content + + +def test_parse_plugin_outs_strips_remote_version_for_lookup() -> None: + """`remote:` matching is version-tolerant — pinned and unpinned both land in the + caller-supplied suffixes dict.""" + suffixes = {"remote:buf.build/protocolbuffers/python": "_pb2"} + pinned = dedent( + """\ + version: v2 + plugins: + - remote: buf.build/protocolbuffers/python:v34.1 + revision: 1 + out: gen + """ + ).encode("utf-8") + unpinned = dedent( + """\ + version: v2 + plugins: + - remote: buf.build/protocolbuffers/python + out: gen + """ + ).encode("utf-8") + assert parse_plugin_outs(pinned, suffixes) == {"_pb2": "gen"} + assert parse_plugin_outs(unpinned, suffixes) == {"_pb2": "gen"} + + +def test_parse_buf_yaml_deps_extracts_module_ids() -> None: + content = dedent( + """\ + version: v2 + modules: + - path: idl + deps: + - buf.build/bufbuild/protovalidate + - buf.build/googleapis/googleapis + """ + ).encode("utf-8") + assert parse_buf_yaml_deps(content) == ( + "buf.build/bufbuild/protovalidate", + "buf.build/googleapis/googleapis", + ) + + +def test_parse_buf_yaml_deps_returns_empty_for_missing_or_invalid() -> None: + no_deps = dedent("version: v2\nmodules:\n - path: idl\n").encode("utf-8") + assert parse_buf_yaml_deps(no_deps) == () + assert parse_buf_yaml_deps(b"not: valid: yaml: ::\nx") == () + assert parse_buf_yaml_deps(b"") == () + + +def test_suffix_plugin_includes_imports_true_for_remote() -> None: + content = dedent( + """\ + version: v2 + plugins: + - remote: buf.build/protocolbuffers/python + out: gen + include_imports: true + """ + ).encode("utf-8") + suffixes = {"remote:buf.build/protocolbuffers/python": "_pb2"} + assert suffix_plugin_includes_imports(content, "_pb2", suffixes) is True + + +def test_suffix_plugin_includes_imports_false_when_unset() -> None: + content = dedent( + """\ + version: v2 + plugins: + - remote: buf.build/protocolbuffers/python + out: gen + """ + ).encode("utf-8") + suffixes = {"remote:buf.build/protocolbuffers/python": "_pb2"} + assert suffix_plugin_includes_imports(content, "_pb2", suffixes) is False + + +def test_suffix_plugin_includes_imports_only_checks_matching_suffix() -> None: + """`include_imports` on a different-suffix plugin doesn't bleed over.""" + content = dedent( + """\ + version: v2 + plugins: + - remote: buf.build/protocolbuffers/pyi + out: gen + include_imports: true + - remote: buf.build/protocolbuffers/python + out: gen + """ + ).encode("utf-8") + suffixes = { + "remote:buf.build/protocolbuffers/pyi": "_pb2.pyi", + "remote:buf.build/protocolbuffers/python": "_pb2", + } + assert suffix_plugin_includes_imports(content, "_pb2", suffixes) is False + + +def test_suffix_plugin_includes_imports_tolerates_pinned_remote() -> None: + content = dedent( + """\ + version: v2 + plugins: + - remote: buf.build/protocolbuffers/python:v34.1 + revision: 1 + out: gen + include_imports: true + """ + ).encode("utf-8") + suffixes = {"remote:buf.build/protocolbuffers/python": "_pb2"} + assert suffix_plugin_includes_imports(content, "_pb2", suffixes) is True + + +def test_suffix_plugin_includes_imports_works_for_protoc_builtin() -> None: + """`include_imports` is a buf-level switch, applicable to `protoc_builtin:` + plugins identically — not just `remote:` ones.""" + content = dedent( + """\ + version: v2 + plugins: + - protoc_builtin: python + out: gen + include_imports: true + """ + ).encode("utf-8") + suffixes = {"protoc_builtin:python": "_pb2"} + assert suffix_plugin_includes_imports(content, "_pb2", suffixes) is True diff --git a/src/python/pants/backend/codegen/protobuf/buf/fields.py b/src/python/pants/backend/codegen/protobuf/buf/fields.py new file mode 100644 index 00000000000..08a200266b4 --- /dev/null +++ b/src/python/pants/backend/codegen/protobuf/buf/fields.py @@ -0,0 +1,41 @@ +# Copyright 2026 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +"""Language-agnostic per-target buf fields. + +`BufGenTemplateField` is shared across language backends (Python, Go, JS, ...) — a +single `buf.gen.yaml` typically declares plugins for multiple languages. +""" + +from __future__ import annotations + +from pants.backend.codegen.protobuf.target_types import ( + ProtobufSourcesGeneratorTarget, + ProtobufSourceTarget, +) +from pants.engine.target import StringField +from pants.util.strutil import help_text + + +class BufGenTemplateField(StringField): + alias = "buf_gen_template" + default = None + help = help_text( + """ + Path to a `buf.gen.yaml` template for this target, overriding + `[buf].gen_template`. The path is interpreted relative to the BUILD file's + directory. + + Only consulted when the target opts into buf-based code generation + via `protobuf_generator='buf'`. + """ + ) + + +def rules(): + return [ + ProtobufSourceTarget.register_plugin_field(BufGenTemplateField), + ProtobufSourcesGeneratorTarget.register_plugin_field( + BufGenTemplateField, as_moved_field=True + ), + ] diff --git a/src/python/pants/backend/codegen/protobuf/buf/lockfile.py b/src/python/pants/backend/codegen/protobuf/buf/lockfile.py new file mode 100644 index 00000000000..6db6c386c1a --- /dev/null +++ b/src/python/pants/backend/codegen/protobuf/buf/lockfile.py @@ -0,0 +1,137 @@ +# Copyright 2026 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +"""Generate `buf.lock` for buf-managed proto modules. + +Each `buf.yaml` in the repo defines a separate buf module (resolve). Running +`pants generate-lockfiles` invokes `buf dep update` for each, producing a +fully-pinned `buf.lock` next to the corresponding `buf.yaml`. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + +from pants.backend.codegen.protobuf.buf.subsystem import BufSubsystem +from pants.core.goals.generate_lockfiles import ( + GenerateLockfile, + GenerateLockfileResult, + KnownUserResolveNames, + KnownUserResolveNamesRequest, + RequestedUserResolveNames, + UserGenerateLockfiles, +) +from pants.core.util_rules.config_files import find_config_file +from pants.core.util_rules.external_tool import download_external_tool +from pants.engine.fs import MergeDigests, PathGlobs +from pants.engine.intrinsics import merge_digests, path_globs_to_digest +from pants.engine.platform import Platform +from pants.engine.process import Process, execute_process_or_raise +from pants.engine.rules import collect_rules, concurrently, implicitly, rule +from pants.engine.unions import UnionRule +from pants.util.logging import LogLevel + + +class KnownBufResolveNamesRequest(KnownUserResolveNamesRequest): + pass + + +class RequestedBufResolveNames(RequestedUserResolveNames): + pass + + +@dataclass(frozen=True) +class GenerateBufLockfile(GenerateLockfile): + """A request to (re)generate `buf.lock` for a single buf module.""" + + buf_yaml_path: str + + +def _resolve_name(buf_yaml_path: str) -> str: + parent = os.path.dirname(buf_yaml_path) + return parent if parent else "buf" + + +@rule +async def known_buf_user_resolve_names( + _: KnownBufResolveNamesRequest, buf: BufSubsystem +) -> KnownUserResolveNames: + files = await find_config_file(buf.config_request) + yaml_paths = sorted(p for p in files.snapshot.files if os.path.basename(p) == "buf.yaml") + return KnownUserResolveNames( + names=tuple(_resolve_name(p) for p in yaml_paths), + option_name="`buf.yaml` discovery", + requested_resolve_names_cls=RequestedBufResolveNames, + ) + + +@rule +async def setup_user_buf_lockfile_requests( + requested: RequestedBufResolveNames, buf: BufSubsystem +) -> UserGenerateLockfiles: + files = await find_config_file(buf.config_request) + name_to_yaml = { + _resolve_name(p): p for p in files.snapshot.files if os.path.basename(p) == "buf.yaml" + } + requests = [] + for name in requested: + path = name_to_yaml.get(name) + if path is None: + continue + requests.append( + GenerateBufLockfile( + resolve_name=name, + lockfile_dest=os.path.join(os.path.dirname(path), "buf.lock"), + diff=False, + buf_yaml_path=path, + ) + ) + return UserGenerateLockfiles(requests) + + +@rule(desc="Resolve buf.yaml deps via `buf dep update`", level=LogLevel.DEBUG) +async def generate_buf_lockfile( + req: GenerateBufLockfile, buf: BufSubsystem, platform: Platform +) -> GenerateLockfileResult: + buf_yaml_dir = os.path.dirname(req.buf_yaml_path) or "." + + # `buf dep update` does a partial build, so it needs at least one `.proto` + # file in the module to operate. Glob them in alongside the buf.yaml/buf.lock. + proto_glob = f"{buf_yaml_dir}/**/*.proto" if buf_yaml_dir != "." else "**/*.proto" + downloaded_buf, files, protos_digest = await concurrently( + download_external_tool(buf.get_request(platform)), + find_config_file(buf.config_request), + path_globs_to_digest(PathGlobs([proto_glob])), + ) + + input_digest = await merge_digests( + MergeDigests((files.snapshot.digest, protos_digest, downloaded_buf.digest)) + ) + + process_result = await execute_process_or_raise( + **implicitly( + Process( + argv=[downloaded_buf.exe, "dep", "update", buf_yaml_dir], + input_digest=input_digest, + description=f"Resolving buf.lock for `{req.resolve_name}`", + level=LogLevel.DEBUG, + output_files=(req.lockfile_dest,), + ) + ) + ) + + return GenerateLockfileResult( + digest=process_result.output_digest, + resolve_name=req.resolve_name, + path=req.lockfile_dest, + ) + + +def rules(): + return [ + *collect_rules(), + UnionRule(GenerateLockfile, GenerateBufLockfile), + UnionRule(KnownUserResolveNamesRequest, KnownBufResolveNamesRequest), + UnionRule(RequestedUserResolveNames, RequestedBufResolveNames), + ] diff --git a/src/python/pants/backend/codegen/protobuf/buf/lockfile_test.py b/src/python/pants/backend/codegen/protobuf/buf/lockfile_test.py new file mode 100644 index 00000000000..f197e755e35 --- /dev/null +++ b/src/python/pants/backend/codegen/protobuf/buf/lockfile_test.py @@ -0,0 +1,63 @@ +# Copyright 2026 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +from __future__ import annotations + +import pytest + +from pants.backend.codegen.protobuf.buf.lockfile import ( + GenerateBufLockfile, + KnownBufResolveNamesRequest, + RequestedBufResolveNames, + _resolve_name, +) +from pants.backend.codegen.protobuf.buf.lockfile import ( + rules as lockfile_rules, +) +from pants.core.goals.generate_lockfiles import KnownUserResolveNames, UserGenerateLockfiles +from pants.testutil.rule_runner import QueryRule, RuleRunner + + +@pytest.fixture +def rule_runner() -> RuleRunner: + return RuleRunner( + rules=[ + *lockfile_rules(), + QueryRule(KnownUserResolveNames, [KnownBufResolveNamesRequest]), + QueryRule(UserGenerateLockfiles, [RequestedBufResolveNames]), + ], + ) + + +def test_resolve_name_uses_parent_directory_or_buf_for_root() -> None: + assert _resolve_name("buf.yaml") == "buf" + assert _resolve_name("idl/buf.yaml") == "idl" + assert _resolve_name("a/b/c/buf.yaml") == "a/b/c" + + +def test_known_resolve_names_finds_repo_root_buf_yaml(rule_runner: RuleRunner) -> None: + rule_runner.write_files({"buf.yaml": "version: v2\nmodules:\n - path: .\n"}) + result = rule_runner.request(KnownUserResolveNames, [KnownBufResolveNamesRequest()]) + assert result.names == ("buf",) + assert result.requested_resolve_names_cls is RequestedBufResolveNames + + +def test_known_resolve_names_returns_empty_when_no_buf_yaml(rule_runner: RuleRunner) -> None: + result = rule_runner.request(KnownUserResolveNames, [KnownBufResolveNamesRequest()]) + assert result.names == () + + +def test_setup_lockfile_requests_maps_resolve_name_to_buf_yaml(rule_runner: RuleRunner) -> None: + rule_runner.write_files({"buf.yaml": "version: v2\nmodules:\n - path: .\n"}) + result = rule_runner.request(UserGenerateLockfiles, [RequestedBufResolveNames(["buf"])]) + [req] = result + assert isinstance(req, GenerateBufLockfile) + assert req.resolve_name == "buf" + assert req.buf_yaml_path == "buf.yaml" + assert req.lockfile_dest == "buf.lock" + + +def test_setup_lockfile_requests_skips_unknown_resolve_names(rule_runner: RuleRunner) -> None: + rule_runner.write_files({"buf.yaml": "version: v2\nmodules:\n - path: .\n"}) + result = rule_runner.request(UserGenerateLockfiles, [RequestedBufResolveNames(["nonexistent"])]) + assert list(result) == [] diff --git a/src/python/pants/backend/codegen/protobuf/lint/buf/skip_field.py b/src/python/pants/backend/codegen/protobuf/buf/skip_field.py similarity index 100% rename from src/python/pants/backend/codegen/protobuf/lint/buf/skip_field.py rename to src/python/pants/backend/codegen/protobuf/buf/skip_field.py diff --git a/src/python/pants/backend/codegen/protobuf/buf/subsystem.py b/src/python/pants/backend/codegen/protobuf/buf/subsystem.py new file mode 100644 index 00000000000..efe9cb889ed --- /dev/null +++ b/src/python/pants/backend/codegen/protobuf/buf/subsystem.py @@ -0,0 +1,155 @@ +# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +from __future__ import annotations + +from pants.core.util_rules.config_files import ConfigFilesRequest +from pants.core.util_rules.external_tool import TemplatedExternalTool +from pants.engine.platform import Platform +from pants.option.option_types import ( + ArgsListOption, + BoolOption, + DictOption, + FileOption, + SkipOption, +) +from pants.util.strutil import softwrap + + +class BufSubsystem(TemplatedExternalTool): + options_scope = "buf" + name = "Buf" + help = "A code generator, linter and formatter for Protocol Buffers (https://github.com/bufbuild/buf)." + + default_version = "v1.69.0" + default_known_versions = [ + "v1.69.0|linux_arm64 |28b258cb4ee7a1224a61e1dd91ae5935f1c86c23e8e67bcfa23c8b096b0ad478|33862042", + "v1.69.0|linux_x86_64|2b1f9cfb5e17d50c10dea9202979ffd28ca7ff7a6f4e51e801a9463986690b03|37497898", + "v1.69.0|macos_arm64 |246534674239f326dd1ad2642e57865dd56dc4e98c4857d13da6eca4d3a168ad|35754114", + "v1.69.0|macos_x86_64|570458723a1e400d654e916cc54a98e78a9c23d72775afde9ef12abcdebd37a1|38439569", + ] + default_url_template = ( + "https://github.com/bufbuild/buf/releases/download/{version}/buf-{platform}.tar.gz" + ) + default_url_platform_mapping = { + "macos_arm64": "Darwin-arm64", + "macos_x86_64": "Darwin-x86_64", + "linux_arm64": "Linux-aarch64", + "linux_x86_64": "Linux-x86_64", + } + + format_skip = SkipOption("fmt", "lint") + lint_skip = SkipOption("lint") + format_args = ArgsListOption(example="--error-format json") + lint_args = ArgsListOption(example="--error-format json") + gen_args = ArgsListOption(example="--include-imports") + + config = FileOption( + default=None, + advanced=True, + help=lambda cls: softwrap( + f""" + Path to a config file understood by Buf + (https://docs.buf.build/configuration/overview). + + Setting this option will disable `[{cls.options_scope}].config_discovery`. Use + this option if the config is located in a non-standard location. + """ + ), + ) + config_discovery = BoolOption( + default=True, + advanced=True, + help=lambda cls: softwrap( + f""" + If true, Pants will include any relevant root config files during runs + (`buf.yaml`). If the json format is preferred, the path to the `buf.json` + file should be provided in the config option. + + Use `[{cls.options_scope}].config` instead if your config is in a non-standard location. + """ + ), + ) + + gen_template = FileOption( + default=None, + advanced=True, + help=lambda cls: softwrap( + f""" + Path to a `buf.gen.yaml` template used by `buf generate` + (https://buf.build/docs/configuration/v2/buf-gen-yaml). + + Used when a `protobuf_source` target opts into buf-based code generation + via the `protobuf_generator` field. May be overridden on a per-target + basis via the `buf_gen_template` field. + + Setting this option will disable `[{cls.options_scope}].gen_template_discovery`. Use + this option if the template is located in a non-standard location. + """ + ), + ) + gen_template_discovery = BoolOption( + default=True, + advanced=True, + help=lambda cls: softwrap( + f""" + If true, Pants will include any relevant `buf.gen.yaml` files found at the + repository root during code generation runs. + + Use `[{cls.options_scope}].gen_template` instead if your template is in a + non-standard location, or use the `buf_gen_template` field for a per-target + override. + """ + ), + ) + + extra_plugin_pins = DictOption[str]( + default={}, + help=softwrap( + """ + Map of `buf.gen.yaml` plugin ids to default `vX.Y:revN` pins that + Pants will fill in for unpinned `remote:` entries, layered on top of + Pants's built-in registry. + + Pants requires every `remote:` plugin to be pinned to an exact + version+revision so codegen output is reproducible. For plugins in + Pants's built-in registry, you can omit the pin and Pants synthesizes + it. Use this option to do the same for custom or forked plugins, or + to override the built-in default for a known plugin. + + Example: + + extra_plugin_pins = { + "myorg.example.com/internal/python-fork": "v2.0:3", + } + """ + ), + advanced=True, + ) + + @property + def config_request(self) -> ConfigFilesRequest: + # Refer to https://docs.buf.build/configuration/overview. + # `buf.lock` is included so codegen / inference invalidate when BSR + # `deps:` resolve to different module versions. + return ConfigFilesRequest( + specified=self.config, + specified_option_name=f"{self.options_scope}.config", + discovery=self.config_discovery, + check_existence=("buf.yaml", "buf.lock"), + check_content={}, + ) + + @property + def gen_template_request(self) -> ConfigFilesRequest: + # Refer to https://buf.build/docs/configuration/v2/buf-gen-yaml. + return ConfigFilesRequest( + specified=self.gen_template, + specified_option_name=f"{self.options_scope}.gen_template", + discovery=self.gen_template_discovery, + check_existence=("buf.gen.yaml",), + check_content={}, + ) + + def generate_exe(self, plat: Platform) -> str: + return "./buf/bin/buf" diff --git a/src/python/pants/backend/codegen/protobuf/lint/buf/format_rules.py b/src/python/pants/backend/codegen/protobuf/lint/buf/format_rules.py index 7b468c79154..8f6f6143e80 100644 --- a/src/python/pants/backend/codegen/protobuf/lint/buf/format_rules.py +++ b/src/python/pants/backend/codegen/protobuf/lint/buf/format_rules.py @@ -2,8 +2,8 @@ # Licensed under the Apache License, Version 2.0 (see LICENSE). from dataclasses import dataclass -from pants.backend.codegen.protobuf.lint.buf.skip_field import SkipBufFormatField -from pants.backend.codegen.protobuf.lint.buf.subsystem import BufSubsystem +from pants.backend.codegen.protobuf.buf.skip_field import SkipBufFormatField +from pants.backend.codegen.protobuf.buf.subsystem import BufSubsystem from pants.backend.codegen.protobuf.target_types import ( ProtobufDependenciesField, ProtobufSourceField, diff --git a/src/python/pants/backend/codegen/protobuf/lint/buf/lint_rules.py b/src/python/pants/backend/codegen/protobuf/lint/buf/lint_rules.py index 4e26b8db1b1..1d24cdb2e1a 100644 --- a/src/python/pants/backend/codegen/protobuf/lint/buf/lint_rules.py +++ b/src/python/pants/backend/codegen/protobuf/lint/buf/lint_rules.py @@ -3,8 +3,8 @@ from dataclasses import dataclass from typing import Any -from pants.backend.codegen.protobuf.lint.buf.skip_field import SkipBufLintField -from pants.backend.codegen.protobuf.lint.buf.subsystem import BufSubsystem +from pants.backend.codegen.protobuf.buf.skip_field import SkipBufLintField +from pants.backend.codegen.protobuf.buf.subsystem import BufSubsystem from pants.backend.codegen.protobuf.target_types import ( ProtobufDependenciesField, ProtobufSourceField, diff --git a/src/python/pants/backend/codegen/protobuf/lint/buf/register.py b/src/python/pants/backend/codegen/protobuf/lint/buf/register.py index c9d216ea0af..c28edc50fb1 100644 --- a/src/python/pants/backend/codegen/protobuf/lint/buf/register.py +++ b/src/python/pants/backend/codegen/protobuf/lint/buf/register.py @@ -1,10 +1,11 @@ # Copyright 2022 Pants project contributors (see CONTRIBUTORS.md). # Licensed under the Apache License, Version 2.0 (see LICENSE). -from pants.backend.codegen.protobuf.lint.buf import skip_field +from pants.backend.codegen.protobuf.buf import skip_field +from pants.backend.codegen.protobuf.buf.lockfile import rules as buf_lockfile_rules +from pants.backend.codegen.protobuf.buf.subsystem import BufSubsystem from pants.backend.codegen.protobuf.lint.buf.format_rules import rules as buf_format_rules from pants.backend.codegen.protobuf.lint.buf.lint_rules import rules as buf_lint_rules -from pants.backend.codegen.protobuf.lint.buf.subsystem import BufSubsystem from pants.core.goals.resolves import ExportableTool from pants.engine.unions import UnionRule @@ -13,6 +14,7 @@ def rules(): return ( *buf_format_rules(), *buf_lint_rules(), + *buf_lockfile_rules(), *skip_field.rules(), UnionRule(ExportableTool, BufSubsystem), ) diff --git a/src/python/pants/backend/codegen/protobuf/lint/buf/subsystem.py b/src/python/pants/backend/codegen/protobuf/lint/buf/subsystem.py deleted file mode 100644 index 63b676e46bb..00000000000 --- a/src/python/pants/backend/codegen/protobuf/lint/buf/subsystem.py +++ /dev/null @@ -1,79 +0,0 @@ -# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md). -# Licensed under the Apache License, Version 2.0 (see LICENSE). - -from __future__ import annotations - -from pants.core.util_rules.config_files import ConfigFilesRequest -from pants.core.util_rules.external_tool import TemplatedExternalTool -from pants.engine.platform import Platform -from pants.option.option_types import ArgsListOption, BoolOption, FileOption, SkipOption -from pants.util.strutil import softwrap - - -class BufSubsystem(TemplatedExternalTool): - options_scope = "buf" - name = "Buf" - help = "A linter and formatter for Protocol Buffers (https://github.com/bufbuild/buf)." - - default_version = "v1.3.0" - default_known_versions = [ - "v1.3.0|linux_arm64 |fbfd53c501451b36900247734bfa4cbe86ae05d0f51bc298de8711d5ee374ee5|13940828", - "v1.3.0|linux_x86_64|e29c4283b1cd68ada41fa493171c41d7605750d258fcd6ecdf692a63fae95213|15267162", - "v1.3.0|macos_arm64 |147985d7f2816a545792e38b26178ff4027bf16cd3712f6e387a4e3692a16deb|15391890", - "v1.3.0|macos_x86_64|3b6bd2e5a5dd758178aee01fb067261baf5d31bfebe93336915bfdf7b21928c4|15955291", - ] - default_url_template = ( - "https://github.com/bufbuild/buf/releases/download/{version}/buf-{platform}.tar.gz" - ) - default_url_platform_mapping = { - "macos_arm64": "Darwin-arm64", - "macos_x86_64": "Darwin-x86_64", - "linux_arm64": "Linux-aarch64", - "linux_x86_64": "Linux-x86_64", - } - - format_skip = SkipOption("fmt", "lint") - lint_skip = SkipOption("lint") - format_args = ArgsListOption(example="--error-format json") - lint_args = ArgsListOption(example="--error-format json") - - config = FileOption( - default=None, - advanced=True, - help=lambda cls: softwrap( - f""" - Path to a config file understood by Buf - (https://docs.buf.build/configuration/overview). - - Setting this option will disable `[{cls.options_scope}].config_discovery`. Use - this option if the config is located in a non-standard location. - """ - ), - ) - config_discovery = BoolOption( - default=True, - advanced=True, - help=lambda cls: softwrap( - f""" - If true, Pants will include any relevant root config files during runs - (`buf.yaml`). If the json format is preferred, the path to the `buf.json` - file should be provided in the config option. - - Use `[{cls.options_scope}].config` instead if your config is in a non-standard location. - """ - ), - ) - - @property - def config_request(self) -> ConfigFilesRequest: - # Refer to https://docs.buf.build/configuration/overview. - return ConfigFilesRequest( - specified=self.config, - specified_option_name=f"{self.options_scope}.config", - discovery=self.config_discovery, - check_existence=("buf.yaml",), - check_content={}, - ) - - def generate_exe(self, plat: Platform) -> str: - return "./buf/bin/buf" diff --git a/src/python/pants/backend/codegen/protobuf/python/BUILD b/src/python/pants/backend/codegen/protobuf/python/BUILD index b556ba09a14..a4873e19f4f 100644 --- a/src/python/pants/backend/codegen/protobuf/python/BUILD +++ b/src/python/pants/backend/codegen/protobuf/python/BUILD @@ -21,3 +21,9 @@ python_tests( # We want to make sure the default lockfile for MyPy Protobuf works for both macOS and Linux. tags=["platform_specific_behavior"], ) +python_tests( + name="buf_rules_integration_test", + sources=["buf_rules_integration_test.py"], + timeout=330, + tags=["platform_specific_behavior"], +) diff --git a/src/python/pants/backend/codegen/protobuf/python/additional_fields.py b/src/python/pants/backend/codegen/protobuf/python/additional_fields.py index 3af6d30017e..ee8026c74af 100644 --- a/src/python/pants/backend/codegen/protobuf/python/additional_fields.py +++ b/src/python/pants/backend/codegen/protobuf/python/additional_fields.py @@ -25,6 +25,9 @@ class PythonSourceRootField(StringField): The source root to generate Python sources under. If unspecified, the source root the `protobuf_sources` is under will be used. + + Ignored when `protobuf_generator='buf'`; in that case the output location + is dictated by the `out:` field of the `buf.gen.yaml` template. """ ) diff --git a/src/python/pants/backend/codegen/protobuf/python/buf_rules.py b/src/python/pants/backend/codegen/protobuf/python/buf_rules.py new file mode 100644 index 00000000000..db65c6c24c4 --- /dev/null +++ b/src/python/pants/backend/codegen/protobuf/python/buf_rules.py @@ -0,0 +1,221 @@ +# Copyright 2026 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass + +from pants.backend.codegen.protobuf.buf.config import ( + MissingBufLockError, + gen_template_request_for_target, + parse_buf_yaml_deps, + resolved_template_path, + synthesize_pinned_buf_gen_yaml, +) +from pants.backend.codegen.protobuf.buf.subsystem import BufSubsystem +from pants.backend.codegen.protobuf.protoc import Protoc +from pants.backend.codegen.protobuf.python.additional_fields import PythonSourceRootField +from pants.backend.codegen.protobuf.target_types import ProtobufSourceField +from pants.core.util_rules.config_files import find_config_file +from pants.core.util_rules.external_tool import download_external_tool +from pants.core.util_rules.source_files import SourceFilesRequest, determine_source_files +from pants.engine.fs import CreateDigest, Directory, FileContent, MergeDigests, RemovePrefix +from pants.engine.internals.graph import transitive_targets as transitive_targets_get +from pants.engine.intrinsics import ( + create_digest, + digest_to_snapshot, + get_digest_contents, + merge_digests, + remove_prefix, +) +from pants.engine.platform import Platform +from pants.engine.process import Process, execute_process_or_raise +from pants.engine.rules import collect_rules, concurrently, implicitly, rule +from pants.engine.target import GeneratedSources, Target, TransitiveTargetsRequest +from pants.util.logging import LogLevel + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class GeneratePythonFromProtobufViaBufRequest: + protocol_target: Target + + +@rule(desc="Generate Python from Protobuf via `buf generate`", level=LogLevel.DEBUG) +async def generate_python_from_protobuf_via_buf( + request: GeneratePythonFromProtobufViaBufRequest, + buf: BufSubsystem, + protoc: Protoc, + platform: Platform, +) -> GeneratedSources: + target = request.protocol_target + + if target.get(PythonSourceRootField).value is not None: + logger.warning( + "`python_source_root` is set on %s but `protobuf_generator='buf'`; " + "the field is ignored — output paths come from the `out:` field of " + "`buf.gen.yaml`.", + target.address, + ) + + output_dir = "_generated_files" + create_output_dir_request = create_digest(CreateDigest([Directory(output_dir)])) + + # Buf needs all transitive `.proto` sources to resolve imports, even though only + # the target's own files are passed via `--path`. + transitive_targets = await transitive_targets_get( + TransitiveTargetsRequest([target.address]), **implicitly() + ) + + # Unlike the protoc path, buf operates on original (unstripped) paths because + # the buf module root is determined by `buf.yaml`'s location, not by Pants source + # roots. + all_sources_request = determine_source_files( + SourceFilesRequest( + tgt[ProtobufSourceField] + for tgt in transitive_targets.closure + if tgt.has_field(ProtobufSourceField) + ) + ) + target_sources_request = determine_source_files( + SourceFilesRequest([target[ProtobufSourceField]]) + ) + + download_buf_request = download_external_tool(buf.get_request(platform)) + download_protoc_request = download_external_tool(protoc.get_request(platform)) + config_files_request = find_config_file(buf.config_request) + gen_template_files_request = find_config_file(gen_template_request_for_target(target, buf)) + + ( + downloaded_buf, + downloaded_protoc, + empty_output_dir, + all_sources, + target_sources, + config_files, + gen_template_files, + ) = await concurrently( + download_buf_request, + download_protoc_request, + create_output_dir_request, + all_sources_request, + target_sources_request, + config_files_request, + gen_template_files_request, + ) + + # If the user's `buf.yaml` declares BSR `deps:`, require a sibling + # `buf.lock` so codegen is reproducible. The lock is what pins each dep to + # an exact commit; without it, buf would resolve to whatever is currently + # latest on the BSR. + config_yaml_paths = [ + p for p in config_files.snapshot.files if os.path.basename(p) == "buf.yaml" + ] + config_lock_paths = { + p for p in config_files.snapshot.files if os.path.basename(p) == "buf.lock" + } + if config_yaml_paths: + yaml_path = config_yaml_paths[0] + config_dcs = await get_digest_contents(config_files.snapshot.digest) + yaml_content = next( + (dc.content for dc in config_dcs if dc.path == yaml_path), + b"", + ) + deps = parse_buf_yaml_deps(yaml_content) + expected_lock = os.path.join(os.path.dirname(yaml_path), "buf.lock") + if deps and expected_lock not in config_lock_paths: + resolve_name = os.path.dirname(yaml_path) or "buf" + raise MissingBufLockError( + f"`{yaml_path}` declares `deps:` ({', '.join(deps)}) but no " + f"`{expected_lock}` was found. Pants requires a `buf.lock` so " + f"BSR deps are pinned and codegen is reproducible.\n\n" + f"Run `pants generate-lockfiles --resolve={resolve_name}` to " + f"create it." + ) + + # Resolve every `remote:` plugin to an exact `:vX.Y:revN` pin before + # invoking buf. Unpinned entries that can be filled in from + # `DEFAULT_PLUGIN_PINS` (or the user's `[buf].extra_plugin_pins`) get a + # default; unknown unpinned entries raise. The resulting yaml is written + # into a fresh digest that replaces the user's `buf.gen.yaml` in the + # sandbox, so buf sees a hermetic, fully-pinned config. + gen_template_digest = gen_template_files.snapshot.digest + if gen_template_files.snapshot.files: + gen_template_path = gen_template_files.snapshot.files[0] + gen_template_dcs = await get_digest_contents(gen_template_digest) + gen_template_content = next( + (dc.content for dc in gen_template_dcs if dc.path == gen_template_path), + b"", + ) + synthesized = synthesize_pinned_buf_gen_yaml( + gen_template_content, + gen_template_path, + extra_pins=buf.extra_plugin_pins, + ) + if synthesized != gen_template_content: + gen_template_digest = await create_digest( + CreateDigest( + [FileContent(gen_template_path, synthesized)], + ) + ) + + input_digest = await merge_digests( + MergeDigests( + ( + all_sources.snapshot.digest, + empty_output_dir, + downloaded_buf.digest, + config_files.snapshot.digest, + gen_template_digest, + ) + ) + ) + + config_arg = ["--config", buf.config] if buf.config else [] + template_path = resolved_template_path(target, buf) + template_arg = ["--template", template_path] if template_path else [] + + argv = [ + downloaded_buf.exe, + "generate", + *config_arg, + *template_arg, + "--output", + output_dir, + *buf.gen_args, + "--path", + ",".join(target_sources.snapshot.files), + ] + + # Expose `protoc` (and any plugin binaries co-located with it) on PATH so + # `buf generate` can resolve `protoc_builtin:` and `local: [protoc]` plugin + # entries. + protoc_relpath = "__protoc" + protoc_bin_dir = os.path.join(protoc_relpath, os.path.dirname(downloaded_protoc.exe)) + + result = await execute_process_or_raise( + **implicitly( + Process( + argv=argv, + input_digest=input_digest, + immutable_input_digests={protoc_relpath: downloaded_protoc.digest}, + env={"PATH": protoc_bin_dir}, + description=f"Generating Python from Protobuf via buf for {target.address}.", + level=LogLevel.DEBUG, + output_directories=(output_dir,), + ) + ), + ) + + # Strip the sandbox `output_dir` prefix; the buf.gen.yaml's `out:` paths land at + # exactly the locations the user declared. + normalized = await remove_prefix(RemovePrefix(result.output_digest, output_dir)) + snapshot = await digest_to_snapshot(normalized) + return GeneratedSources(snapshot) + + +def rules(): + return collect_rules() diff --git a/src/python/pants/backend/codegen/protobuf/python/buf_rules_integration_test.py b/src/python/pants/backend/codegen/protobuf/python/buf_rules_integration_test.py new file mode 100644 index 00000000000..7fd27979ebb --- /dev/null +++ b/src/python/pants/backend/codegen/protobuf/python/buf_rules_integration_test.py @@ -0,0 +1,423 @@ +# Copyright 2026 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +from __future__ import annotations + +from textwrap import dedent + +import pytest + +from pants.backend.codegen.protobuf import protobuf_dependency_inference +from pants.backend.codegen.protobuf.protobuf_dependency_inference import ( + InferProtobufDependencies, + ProtobufDependencyInferenceFieldSet, +) +from pants.backend.codegen.protobuf.python import additional_fields +from pants.backend.codegen.protobuf.python.python_protobuf_subsystem import ( + rules as protobuf_subsystem_rules, +) +from pants.backend.codegen.protobuf.python.register import rules as python_protobuf_backend_rules +from pants.backend.codegen.protobuf.python.rules import GeneratePythonFromProtobufRequest +from pants.backend.codegen.protobuf.python.rules import rules as protobuf_rules +from pants.backend.codegen.protobuf.target_types import ( + ProtobufSourceField, + ProtobufSourcesGeneratorTarget, +) +from pants.backend.codegen.protobuf.target_types import rules as protobuf_target_types_rules +from pants.backend.python import target_types_rules as python_target_types_rules +from pants.backend.python.dependency_inference import module_mapper +from pants.core.target_types import rules as core_target_types_rules +from pants.core.util_rules import stripped_source_files +from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest +from pants.engine.addresses import Address +from pants.engine.target import ( + GeneratedSources, + HydratedSources, + HydrateSourcesRequest, + InferredDependencies, + TransitiveTargets, + TransitiveTargetsRequest, +) +from pants.testutil.rule_runner import QueryRule, RuleRunner + +# A minimal `buf.gen.yaml` template that invokes the python plugin built into protoc. +# `buf` shells out to `protoc` for `protoc_builtin` plugins, so `protoc` must be on PATH. +BUF_GEN_YAML = dedent( + """\ + version: v2 + plugins: + - protoc_builtin: python + out: src/proto + """ +) + +BUF_YAML = dedent( + """\ + version: v2 + modules: + - path: idl/proto + """ +) + +SIMPLE_PROTO = dedent( + """\ + syntax = "proto3"; + package foo; + message Person { + string name = 1; + int32 id = 2; + } + """ +) + + +@pytest.fixture +def rule_runner() -> RuleRunner: + return RuleRunner( + rules=[ + *protobuf_rules(), + *python_protobuf_backend_rules(), + *protobuf_dependency_inference.rules(), + *protobuf_subsystem_rules(), + *additional_fields.rules(), + *protobuf_target_types_rules(), + *python_target_types_rules.rules(), + *stripped_source_files.rules(), + *module_mapper.rules(), + *core_target_types_rules(), + QueryRule(HydratedSources, [HydrateSourcesRequest]), + QueryRule(GeneratedSources, [GeneratePythonFromProtobufRequest]), + QueryRule(InferredDependencies, [InferProtobufDependencies]), + QueryRule(TransitiveTargets, [TransitiveTargetsRequest]), + QueryRule(SourceFiles, [SourceFilesRequest]), + ], + target_types=[ProtobufSourcesGeneratorTarget], + ) + + +def _assert_generates( + rule_runner: RuleRunner, + address: Address, + *, + expected_files: set[str], + source_roots: list[str], +) -> None: + rule_runner.set_options( + [ + f"--source-root-patterns={repr(source_roots)}", + "--no-python-protobuf-infer-runtime-dependency", + ], + env_inherit={"PATH"}, + ) + tgt = rule_runner.get_target(address) + protocol_sources = rule_runner.request( + HydratedSources, [HydrateSourcesRequest(tgt[ProtobufSourceField])] + ) + generated = rule_runner.request( + GeneratedSources, + [GeneratePythonFromProtobufRequest(protocol_sources.snapshot, tgt)], + ) + assert set(generated.snapshot.files) == expected_files + + +@pytest.mark.platform_specific_behavior +def test_buf_generates_python_at_out_directory(rule_runner: RuleRunner) -> None: + """`out: src/proto` in `buf.gen.yaml` lands generated files at `src/proto/...`.""" + rule_runner.write_files( + { + "buf.yaml": BUF_YAML, + "buf.gen.yaml": BUF_GEN_YAML, + "idl/proto/foo/person.proto": SIMPLE_PROTO, + "idl/proto/foo/BUILD": ("protobuf_sources(protobuf_generator='buf')"), + } + ) + _assert_generates( + rule_runner, + Address("idl/proto/foo", relative_file_path="person.proto"), + expected_files={"src/proto/foo/person_pb2.py"}, + source_roots=["idl/proto", "src/proto"], + ) + + +@pytest.mark.platform_specific_behavior +def test_buf_per_target_template_override(rule_runner: RuleRunner) -> None: + """Per-target `buf_gen_template` overrides the global discovered template.""" + rule_runner.write_files( + { + "buf.yaml": BUF_YAML, + # Repo-root template that would write to `gen_default/...` if it were used. + "buf.gen.yaml": dedent( + """\ + version: v2 + plugins: + - protoc_builtin: python + out: gen_default + """ + ), + "idl/proto/foo/buf.gen.yaml": dedent( + """\ + version: v2 + plugins: + - protoc_builtin: python + out: gen_override + """ + ), + "idl/proto/foo/person.proto": SIMPLE_PROTO, + "idl/proto/foo/BUILD": ( + "protobuf_sources(protobuf_generator='buf', buf_gen_template='buf.gen.yaml')" + ), + } + ) + _assert_generates( + rule_runner, + Address("idl/proto/foo", relative_file_path="person.proto"), + expected_files={"gen_override/foo/person_pb2.py"}, + source_roots=["idl/proto", "gen_override"], + ) + + +@pytest.mark.platform_specific_behavior +def test_buf_generates_python_with_remote_plugin(rule_runner: RuleRunner) -> None: + """A `remote:` plugin (fetched from the buf.build registry) is invoked over the + network at codegen time and must produce output the same way `protoc_builtin:` + plugins do. Pinned to a specific version so the test is deterministic.""" + rule_runner.write_files( + { + "buf.yaml": BUF_YAML, + "buf.gen.yaml": dedent( + """\ + version: v2 + plugins: + - remote: buf.build/protocolbuffers/python:v34.1 + out: gen_remote + """ + ), + "idl/proto/foo/person.proto": SIMPLE_PROTO, + "idl/proto/foo/BUILD": "protobuf_sources(protobuf_generator='buf')", + } + ) + _assert_generates( + rule_runner, + Address("idl/proto/foo", relative_file_path="person.proto"), + expected_files={"gen_remote/foo/person_pb2.py"}, + source_roots=["idl/proto", "gen_remote"], + ) + + +@pytest.mark.platform_specific_behavior +def test_buf_only_sends_transitive_closure_to_sandbox(rule_runner: RuleRunner) -> None: + """Codegen for one proto target must (a) include protos it transitively + imports, even from separate folders with separate BUILD files, and (b) + *not* include unrelated siblings in the same buf module. + + Verified end-to-end: + - The target's proto imports a sibling-folder proto. If our `transitive_targets.closure` + missed the import, that sibling wouldn't reach the sandbox and buf would fail to + compile (unresolved import). + - An unrelated sibling proto is deliberately malformed — buf would reject it if it + ever saw it. Codegen succeeding proves it was filtered out. + + Matters for monorepo scale: a buf module with thousands of `.proto` files + must not send all of them to every per-target codegen invocation.""" + rule_runner.write_files( + { + "buf.yaml": BUF_YAML, + "buf.gen.yaml": BUF_GEN_YAML, + # The target's own proto, importing a sibling-folder proto. + "idl/proto/foo/person.proto": dedent( + """\ + syntax = "proto3"; + package foo; + import "common/address.proto"; + message Person { + string name = 1; + common.Address address = 2; + } + """ + ), + "idl/proto/foo/BUILD": "protobuf_sources(protobuf_generator='buf')", + # The dep, in its own folder + BUILD — Pants's proto-dep inference + # picks it up via the `import` statement above. + "idl/proto/common/address.proto": dedent( + """\ + syntax = "proto3"; + package common; + message Address { + string street = 1; + string city = 2; + } + """ + ), + "idl/proto/common/BUILD": "protobuf_sources(protobuf_generator='buf')", + # Unrelated proto in yet another folder, with its own BUILD. Must + # NOT reach the sandbox — deliberately malformed so buf would error + # on it if our closure leaked. + "idl/proto/bar/garbage.proto": "this is not a valid proto file !!!", + "idl/proto/bar/BUILD": "protobuf_sources(protobuf_generator='buf')", + } + ) + _assert_generates( + rule_runner, + Address("idl/proto/foo", relative_file_path="person.proto"), + expected_files={"src/proto/foo/person_pb2.py"}, + source_roots=["idl/proto", "src/proto"], + ) + + # And the dep inferrer reports the same shape: address.proto is in, garbage is out. + person_addr = Address("idl/proto/foo", relative_file_path="person.proto") + person_tgt = rule_runner.get_target(person_addr) + inferred = rule_runner.request( + InferredDependencies, + [InferProtobufDependencies(ProtobufDependencyInferenceFieldSet.create(person_tgt))], + ) + inferred_addrs = set(inferred.include) + assert Address("idl/proto/common", relative_file_path="address.proto") in inferred_addrs + assert Address("idl/proto/bar", relative_file_path="garbage.proto") not in inferred_addrs + + +@pytest.mark.platform_specific_behavior +def test_buf_isolates_per_proto_within_one_build_file(rule_runner: RuleRunner) -> None: + """Cache-invalidation correctness for the case where two protos share a + `protobuf_sources()` glob: per-file targets must remain independent so a + monorepo with thousands of protos under one BUILD doesn't pay a cache tax + on every edit. + + Verifies both: + - **Structural** (the input digest): `garbage.proto` is *not* in + person.proto's `transitive_targets.closure` or its `SourceFiles` digest. + If those bytes aren't part of person's input, no change to garbage can + bust person's cache. + - **E2E**: codegen for `person.proto` succeeds (closure isolates), and + codegen for the malformed `garbage.proto` *fails* — sanity-checking + that the malformed-proto trick has teeth so the negative half of the + structural assertion is meaningful.""" + from pants.engine.internals.scheduler import ExecutionError + + rule_runner.write_files( + { + "buf.yaml": BUF_YAML, + "buf.gen.yaml": BUF_GEN_YAML, + # Both protos under one BUILD file's `protobuf_sources()` glob. + "idl/proto/foo/person.proto": SIMPLE_PROTO, + "idl/proto/foo/garbage.proto": "this is not a valid proto file !!!", + "idl/proto/foo/BUILD": "protobuf_sources(protobuf_generator='buf')", + } + ) + rule_runner.set_options( + [ + "--source-root-patterns=['idl/proto', 'src/proto']", + "--no-python-protobuf-infer-runtime-dependency", + ], + env_inherit={"PATH"}, + ) + + person_addr = Address("idl/proto/foo", relative_file_path="person.proto") + garbage_addr = Address("idl/proto/foo", relative_file_path="garbage.proto") + + # Structural: garbage isn't in person's closure or source-files digest. + transitive = rule_runner.request(TransitiveTargets, [TransitiveTargetsRequest([person_addr])]) + proto_addresses = { + tgt.address for tgt in transitive.closure if tgt.has_field(ProtobufSourceField) + } + assert person_addr in proto_addresses + assert garbage_addr not in proto_addresses + sources = rule_runner.request( + SourceFiles, + [ + SourceFilesRequest( + tgt[ProtobufSourceField] + for tgt in transitive.closure + if tgt.has_field(ProtobufSourceField) + ) + ], + ) + assert "idl/proto/foo/person.proto" in sources.snapshot.files + assert "idl/proto/foo/garbage.proto" not in sources.snapshot.files + + # E2E: person's codegen succeeds; garbage's fails (sanity-check the trick). + person_tgt = rule_runner.get_target(person_addr) + person_hydrated = rule_runner.request( + HydratedSources, [HydrateSourcesRequest(person_tgt[ProtobufSourceField])] + ) + generated = rule_runner.request( + GeneratedSources, + [GeneratePythonFromProtobufRequest(person_hydrated.snapshot, person_tgt)], + ) + assert set(generated.snapshot.files) == {"src/proto/foo/person_pb2.py"} + + garbage_tgt = rule_runner.get_target(garbage_addr) + garbage_hydrated = rule_runner.request( + HydratedSources, [HydrateSourcesRequest(garbage_tgt[ProtobufSourceField])] + ) + with pytest.raises(ExecutionError): + rule_runner.request( + GeneratedSources, + [GeneratePythonFromProtobufRequest(garbage_hydrated.snapshot, garbage_tgt)], + ) + + +@pytest.mark.platform_specific_behavior +def test_buf_codegen_fails_without_buf_lock_when_deps_declared( + rule_runner: RuleRunner, +) -> None: + """`buf.yaml` declaring `deps:` without a sibling `buf.lock` is rejected at + codegen time with a friendly error pointing at `pants generate-lockfiles + --resolve=…`.""" + from pants.backend.codegen.protobuf.buf.config import MissingBufLockError + from pants.engine.internals.scheduler import ExecutionError + + rule_runner.write_files( + { + "buf.yaml": dedent( + """\ + version: v2 + modules: + - path: idl/proto + deps: + - buf.build/bufbuild/protovalidate + """ + ), + "buf.gen.yaml": BUF_GEN_YAML, + # No buf.lock — codegen must reject this before invoking buf. + "idl/proto/foo/person.proto": SIMPLE_PROTO, + "idl/proto/foo/BUILD": "protobuf_sources(protobuf_generator='buf')", + } + ) + rule_runner.set_options( + [ + "--source-root-patterns=['idl/proto', 'src/proto']", + "--no-python-protobuf-infer-runtime-dependency", + ], + env_inherit={"PATH"}, + ) + tgt = rule_runner.get_target(Address("idl/proto/foo", relative_file_path="person.proto")) + protocol_sources = rule_runner.request( + HydratedSources, [HydrateSourcesRequest(tgt[ProtobufSourceField])] + ) + with pytest.raises(ExecutionError) as excinfo: + rule_runner.request( + GeneratedSources, + [GeneratePythonFromProtobufRequest(protocol_sources.snapshot, tgt)], + ) + cause_types = {type(c).__name__ for c in excinfo.value.wrapped_exceptions} + assert MissingBufLockError.__name__ in cause_types + msg = str(excinfo.value) + assert "buf.build/bufbuild/protovalidate" in msg + assert "pants generate-lockfiles --resolve=" in msg + + +@pytest.mark.platform_specific_behavior +def test_default_protoc_path_still_works(rule_runner: RuleRunner) -> None: + """Regression: `protobuf_generator` unset (default `protoc`) is unchanged.""" + rule_runner.write_files( + { + "src/protobuf/foo/person.proto": SIMPLE_PROTO, + "src/protobuf/foo/BUILD": "protobuf_sources()", + } + ) + _assert_generates( + rule_runner, + Address("src/protobuf/foo", relative_file_path="person.proto"), + expected_files={"src/protobuf/foo/person_pb2.py"}, + source_roots=["src/protobuf"], + ) diff --git a/src/python/pants/backend/codegen/protobuf/python/python_protobuf_dependency_inference_test.py b/src/python/pants/backend/codegen/protobuf/python/python_protobuf_dependency_inference_test.py index 407f2475d8c..687d8455cac 100644 --- a/src/python/pants/backend/codegen/protobuf/python/python_protobuf_dependency_inference_test.py +++ b/src/python/pants/backend/codegen/protobuf/python/python_protobuf_dependency_inference_test.py @@ -6,6 +6,7 @@ from pants.backend.codegen.protobuf import ( protobuf_dependency_inference, ) +from pants.backend.codegen.protobuf.buf import fields as buf_fields from pants.backend.codegen.protobuf.protobuf_dependency_inference import ( InferProtobufDependencies, ProtobufDependencyInferenceFieldSet, @@ -30,6 +31,7 @@ def rule_runner_with_python_resolves() -> RuleRunner: *protobuf_dependency_inference.rules(), *target_types_rules(), *additional_fields.rules(), + *buf_fields.rules(), *python_target_types_rules.rules(), QueryRule(ProtobufMapping, []), QueryRule(InferredDependencies, [InferProtobufDependencies]), diff --git a/src/python/pants/backend/codegen/protobuf/python/python_protobuf_module_mapper.py b/src/python/pants/backend/codegen/protobuf/python/python_protobuf_module_mapper.py index 9d4770e3b3e..40ad50a7434 100644 --- a/src/python/pants/backend/codegen/protobuf/python/python_protobuf_module_mapper.py +++ b/src/python/pants/backend/codegen/protobuf/python/python_protobuf_module_mapper.py @@ -3,12 +3,31 @@ from __future__ import annotations +import logging +import os from collections import defaultdict +from collections.abc import Mapping, Sequence +from dataclasses import dataclass from typing import DefaultDict -from pants.backend.codegen.protobuf.python.python_protobuf_subsystem import PythonProtobufSubsystem +from pants.backend.codegen.protobuf.buf.config import ( + BufGenContent, + BufLayout, + fetch_buf_gen_contents, + fetch_buf_layout, + parse_plugin_outs, + suffix_plugin_includes_imports, +) +from pants.backend.codegen.protobuf.buf.subsystem import BufSubsystem +from pants.backend.codegen.protobuf.python.additional_fields import PythonSourceRootField +from pants.backend.codegen.protobuf.python.python_protobuf_subsystem import ( + DEFAULT_BSR_DEP_MODULES, + DEFAULT_PLUGIN_SUFFIXES, + PythonProtobufSubsystem, +) from pants.backend.codegen.protobuf.target_types import ( AllProtobufTargets, + ProtobufGeneratorField, ProtobufGrpcToggleField, ProtobufSourceField, ) @@ -23,9 +42,12 @@ from pants.backend.python.target_types import PythonResolveField from pants.core.util_rules.stripped_source_files import StrippedFileNameRequest, strip_file_name from pants.engine.rules import collect_rules, concurrently, rule +from pants.engine.target import Target from pants.engine.unions import UnionRule from pants.util.logging import LogLevel +logger = logging.getLogger(__name__) + def proto_path_to_py_module(stripped_path: str, *, suffix: str) -> str: return stripped_path.replace(".proto", suffix).replace("/", ".") @@ -36,28 +58,200 @@ class PythonProtobufMappingMarker(FirstPartyPythonMappingImplMarker): pass +# Suffixes relevant to Python codegen. Each is registered iff its plugin appears +# in `buf.gen.yaml`. +_PB2_SUFFIX = "_pb2" +_SERVICE_SUFFIXES: tuple[str, ...] = ("_pb2_grpc", "_grpc", "_connect") + + +@dataclass(frozen=True) +class _BufStripPlan: + """Plan for one buf target: paths to feed to `strip_file_name`, paired with the + suffix to apply to each stripped result.""" + + suffixes: tuple[str, ...] + paths_to_strip: tuple[str, ...] + + +def _plan_buf_target( + target: Target, + suffix_outs: Mapping[str, str], + buf_module_root: str, +) -> _BufStripPlan: + """Build the strip plan from the suffixes matched in this target's `buf.gen.yaml`. + + A suffix is registered iff its plugin appears in the file. `grpc=True` is + *not* consulted: `buf.gen.yaml` is the authoritative source of truth for + which buf target outputs exist. + """ + proto_path = target[ProtobufSourceField].file_path + rel_proto = ( + os.path.relpath(proto_path, buf_module_root) + if buf_module_root + and (proto_path == buf_module_root or proto_path.startswith(buf_module_root + os.sep)) + else proto_path + ) + + def _path_for(out_dir: str) -> str: + return os.path.normpath(os.path.join(out_dir, rel_proto)) + + suffixes: list[str] = [] + paths: list[str] = [] + if _PB2_SUFFIX in suffix_outs: + suffixes.append(_PB2_SUFFIX) + paths.append(_path_for(suffix_outs[_PB2_SUFFIX])) + for suffix in _SERVICE_SUFFIXES: + if suffix in suffix_outs: + suffixes.append(suffix) + paths.append(_path_for(suffix_outs[suffix])) + + return _BufStripPlan(tuple(suffixes), tuple(paths)) + + +def _fallback_plan(target: Target) -> _BufStripPlan: + """Plan when no `buf.gen.yaml` was found: assume the proto's source root also + covers the generated `.py`. Service suffixes can't be inferred without the + template, so we register only `_pb2`.""" + proto_path = target[ProtobufSourceField].file_path + return _BufStripPlan((_PB2_SUFFIX,), (proto_path,)) + + +# Protoc-only subsystem options. Their default values are mirrored here so we can +# detect when the user explicitly set them while also using buf targets, and warn +# that they're ignored on the buf path. Keep in sync with the option definitions +# in `python_protobuf_subsystem.py`. +_PROTOC_ONLY_OPTION_DEFAULTS: tuple[tuple[str, object], ...] = ( + ("grpcio_plugin", True), + ("grpclib_plugin", False), + ("mypy_plugin", False), + ("generate_type_stubs", False), +) + + +def _emit_subsystem_warnings_for_buf(subsystem: PythonProtobufSubsystem) -> None: + """Warn once if subsystem options that are protoc-only are set non-default + while at least one buf target exists.""" + for option_name, default in _PROTOC_ONLY_OPTION_DEFAULTS: + if getattr(subsystem, option_name) == default: + continue + logger.warning( + "[%s].%s is set but ignored for `protobuf_generator='buf'` targets. " + "Service generation and `.pyi` stubs for buf targets are determined by " + "the plugin entries in `buf.gen.yaml`.", + subsystem.options_scope, + option_name, + ) + + +def _emit_per_target_warnings_for_buf(target: Target) -> None: + """Warn once per buf target about field values that have no effect.""" + if target.get(ProtobufGrpcToggleField).value: + logger.warning( + "`grpc=True` is set on %s but is ignored for `protobuf_generator='buf'` " + "targets. Whether `_pb2_grpc.py` / `_grpc.py` / `_connect.py` exist is " + "determined by the plugins in `buf.gen.yaml`.", + target.address, + ) + if target.get(PythonSourceRootField).value is not None: + logger.warning( + "`python_source_root` is set on %s but ignored for " + "`protobuf_generator='buf'`; output paths come from the `out:` field of " + "`buf.gen.yaml`.", + target.address, + ) + + @rule(desc="Creating map of Protobuf targets to generated Python modules", level=LogLevel.DEBUG) async def map_protobuf_to_python_modules( protobuf_targets: AllProtobufTargets, python_setup: PythonSetup, python_protobuf_subsystem: PythonProtobufSubsystem, + buf: BufSubsystem, _: PythonProtobufMappingMarker, ) -> FirstPartyPythonMappingImpl: - grpc_suffixes = [] + grpc_suffixes_list: list[str] = [] if python_protobuf_subsystem.grpcio_plugin: - grpc_suffixes.append("_pb2_grpc") + grpc_suffixes_list.append("_pb2_grpc") if python_protobuf_subsystem.grpclib_plugin: - grpc_suffixes.append("_grpc") + grpc_suffixes_list.append("_grpc") + grpc_suffixes = tuple(grpc_suffixes_list) + + protoc_targets: list[Target] = [] + buf_targets: list[Target] = [] + for tgt in protobuf_targets: + if tgt.get(ProtobufGeneratorField).value == "buf": + buf_targets.append(tgt) + else: + protoc_targets.append(tgt) + + if buf_targets: + _emit_subsystem_warnings_for_buf(python_protobuf_subsystem) - stripped_file_per_target = await concurrently( + # ---- protoc path. ---- + stripped_file_per_protoc_target = await concurrently( strip_file_name(StrippedFileNameRequest(tgt[ProtobufSourceField].file_path)) - for tgt in protobuf_targets + for tgt in protoc_targets + ) + + # ---- buf path: registry-driven plugin matching, no `grpc=True` gate. ---- + if buf_targets: + buf_layout: BufLayout = await fetch_buf_layout(buf) + buf_gen_contents: tuple[BufGenContent, ...] = await fetch_buf_gen_contents(buf_targets, buf) + else: + buf_layout = BufLayout("", (), ()) + buf_gen_contents = () + + plugin_suffixes = { + **DEFAULT_PLUGIN_SUFFIXES, + **python_protobuf_subsystem.extra_buf_plugin_suffixes, + } + plans: list[_BufStripPlan] = [] + for gen in buf_gen_contents: + _emit_per_target_warnings_for_buf(gen.target) + if gen.template_path is None: + logger.debug( + "No `buf.gen.yaml` resolved for %s; falling back to source-root path " + "arithmetic for `_pb2` only. Service suffixes can't be inferred without " + "the template.", + gen.target.address, + ) + plans.append(_fallback_plan(gen.target)) + continue + # Inference doesn't enforce pinning — that's codegen's job. Inference + # works fine on unpinned entries (we only need plugin ids to look up + # suffixes), so the user's editor-side dep inference doesn't fail on + # in-flight `buf.gen.yaml` edits. + suffix_outs = parse_plugin_outs(gen.content, plugin_suffixes) + proto_path = gen.target[ProtobufSourceField].file_path + buf_module_root = buf_layout.root_for_proto(proto_path) + plans.append(_plan_buf_target(gen.target, suffix_outs, buf_module_root)) + + flat_strip_requests: list[StrippedFileNameRequest] = [ + StrippedFileNameRequest(p) for plan in plans for p in plan.paths_to_strip + ] + flat_stripped = ( + await concurrently(strip_file_name(req) for req in flat_strip_requests) + if flat_strip_requests + else () ) + # Reassemble per-target module lists. + buf_modules_per_target: list[list[str]] = [] + idx = 0 + for plan in plans: + target_modules: list[str] = [] + for suffix in plan.suffixes: + stripped = flat_stripped[idx] + target_modules.append(proto_path_to_py_module(stripped.value, suffix=suffix)) + idx += 1 + buf_modules_per_target.append(target_modules) + + # ---- Build the module → providers map. ---- resolves_to_modules_to_providers: DefaultDict[ ResolveName, DefaultDict[str, list[ModuleProvider]] ] = defaultdict(lambda: defaultdict(list)) - for tgt, stripped_file in zip(protobuf_targets, stripped_file_per_target): + + for tgt, stripped_file in zip(protoc_targets, stripped_file_per_protoc_target): resolve = tgt[PythonResolveField].normalized_value(python_setup) # NB: We don't consider the MyPy plugin, which generates `_pb2.pyi`. The stubs end up @@ -74,6 +268,37 @@ async def map_protobuf_to_python_modules( ModuleProvider(tgt.address, ModuleProviderType.IMPL) ) + for tgt, modules in zip(buf_targets, buf_modules_per_target): + resolve = tgt[PythonResolveField].normalized_value(python_setup) + for module in modules: + resolves_to_modules_to_providers[resolve][module].append( + ModuleProvider(tgt.address, ModuleProviderType.IMPL) + ) + + # Register BSR-dep Python modules as owned by buf targets that actually + # generate them. A target only generates BSR-dep `*_pb2.py` files if its + # `buf.gen.yaml` has `include_imports: true` on whichever plugin emits + # `_pb2` (the BSR remote, `protoc_builtin: python`, etc.) — otherwise the + # file isn't in the `GeneratedSources` digest and registering ownership + # would lie. Targets without `include_imports` are skipped; users either set + # it or accept the dep-inference warning. + if buf_layout.deps and buf_targets: + bsr_to_modules: Mapping[str, Sequence[str]] = { + **{k: tuple(v) for k, v in DEFAULT_BSR_DEP_MODULES.items()}, + **{k: tuple(v) for k, v in python_protobuf_subsystem.extra_buf_bsr_modules.items()}, + } + bsr_modules_for_layout: list[str] = [] + for dep in buf_layout.deps: + bsr_modules_for_layout.extend(bsr_to_modules.get(dep, ())) + for tgt, gen in zip(buf_targets, buf_gen_contents): + if not suffix_plugin_includes_imports(gen.content, "_pb2", plugin_suffixes): + continue + resolve = tgt[PythonResolveField].normalized_value(python_setup) + for module in bsr_modules_for_layout: + resolves_to_modules_to_providers[resolve][module].append( + ModuleProvider(tgt.address, ModuleProviderType.IMPL) + ) + return FirstPartyPythonMappingImpl.create(resolves_to_modules_to_providers) diff --git a/src/python/pants/backend/codegen/protobuf/python/python_protobuf_module_mapper_test.py b/src/python/pants/backend/codegen/protobuf/python/python_protobuf_module_mapper_test.py index b4f9f565046..fb3ba0527c6 100644 --- a/src/python/pants/backend/codegen/protobuf/python/python_protobuf_module_mapper_test.py +++ b/src/python/pants/backend/codegen/protobuf/python/python_protobuf_module_mapper_test.py @@ -7,6 +7,7 @@ import pytest +from pants.backend.codegen.protobuf.buf import fields as buf_fields from pants.backend.codegen.protobuf.python import additional_fields, python_protobuf_module_mapper from pants.backend.codegen.protobuf.python.python_protobuf_module_mapper import ( PythonProtobufMappingMarker, @@ -18,7 +19,7 @@ ModuleProvider, ModuleProviderType, ) -from pants.core.util_rules import stripped_source_files +from pants.core.util_rules import config_files, stripped_source_files from pants.engine.addresses import Address from pants.testutil.rule_runner import QueryRule, RuleRunner @@ -27,7 +28,9 @@ def rule_runner() -> RuleRunner: return RuleRunner( rules=[ + *buf_fields.rules(), *additional_fields.rules(), + *config_files.rules(), *stripped_source_files.rules(), *python_protobuf_module_mapper.rules(), *python_protobuf_target_types_rules(), @@ -185,6 +188,459 @@ def providers(addresses: list[Address]) -> tuple[ModuleProvider, ...]: ) +def test_buf_target_falls_back_to_source_root_math_without_gen_yaml( + rule_runner: RuleRunner, +) -> None: + """When `buf.gen.yaml` is absent, the buf path falls back to protoc-style path math. + + This is correct as long as the convention (buf module root and `out:` aligning with + Pants source roots) holds. The test does not write a `buf.gen.yaml`. + """ + rule_runner.set_options(["--source-root-patterns=['src/protobuf']", "--python-enable-resolves"]) + rule_runner.write_files( + { + "src/protobuf/foo/f.proto": "", + "src/protobuf/foo/BUILD": ("protobuf_sources(protobuf_generator='buf')"), + } + ) + result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()]) + assert result == FirstPartyPythonMappingImpl.create( + { + "python-default": { + "foo.f_pb2": ( + ModuleProvider( + Address("src/protobuf/foo", relative_file_path="f.proto"), + ModuleProviderType.IMPL, + ), + ) + } + } + ) + + +def test_buf_target_uses_gen_yaml_out_directory(rule_runner: RuleRunner) -> None: + """When `buf.gen.yaml` declares `out: src/proto`, generated paths sit under that dir. + + Here the `.proto` lives at `idl/proto/foo/bar.proto` (buf module root `idl/proto`), + and `out: src/proto` is the Python source root. Module name is `foo.bar_pb2`. + """ + rule_runner.set_options( + [ + "--source-root-patterns=['idl/proto', 'src/proto']", + "--python-enable-resolves", + ] + ) + rule_runner.write_files( + { + "idl/proto/buf.yaml": "version: v2\nmodules:\n - path: .\n", + "idl/proto/buf.gen.yaml": ( + "version: v2\nplugins:\n - protoc_builtin: python\n out: src/proto\n" + ), + "idl/proto/foo/bar.proto": "", + "idl/proto/foo/BUILD": ("protobuf_sources(protobuf_generator='buf')"), + } + ) + # The discovery glob for `buf.gen.yaml` only checks the repo root, so we stage one + # there too; alternatively a per-target field override would be used. + rule_runner.write_files( + { + "buf.gen.yaml": ( + "version: v2\nplugins:\n - protoc_builtin: python\n out: src/proto\n" + ), + } + ) + result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()]) + assert result == FirstPartyPythonMappingImpl.create( + { + "python-default": { + "foo.bar_pb2": ( + ModuleProvider( + Address("idl/proto/foo", relative_file_path="bar.proto"), + ModuleProviderType.IMPL, + ), + ) + } + } + ) + + +def test_buf_target_default_remote_plugin_matches_with_version_pin( + rule_runner: RuleRunner, +) -> None: + """The default `python_buf_plugin = buf.build/protocolbuffers/python` matches a + `remote:` plugin entry even when the entry includes a `:vXX.X` version suffix.""" + rule_runner.set_options(["--source-root-patterns=['/']", "--python-enable-resolves"]) + rule_runner.write_files( + { + "buf.yaml": "version: v2\nmodules:\n - path: .\n", + "buf.gen.yaml": ( + "version: v2\n" + "plugins:\n" + " - remote: buf.build/protocolbuffers/python:v34.1\n" + " out: gen\n" + ), + "protos/svc.proto": "", + "protos/BUILD": "protobuf_sources(protobuf_generator='buf')", + } + ) + result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()]) + assert result == FirstPartyPythonMappingImpl.create( + { + "python-default": { + "gen.protos.svc_pb2": ( + ModuleProvider( + Address("protos", relative_file_path="svc.proto"), + ModuleProviderType.IMPL, + ), + ) + } + } + ) + + +def test_buf_target_connectrpc_plugin_registers_connect_modules( + rule_runner: RuleRunner, +) -> None: + """A `connectrpc/python` plugin entry registers `_connect.py` modules + automatically — no subsystem opt-in, plugin presence in `buf.gen.yaml` is + sufficient. + + Repo layout (mirrors the connectrpc.com getting-started example): + - `buf.yaml` declares the buf module at `company/proto`. + - `buf.gen.yaml` runs the protobuf-python and connectrpc-python remote plugins, + sending pb2 output to `company/proto/gen` and connect output to `company/protogen`. + - The proto file lives at `company/proto/services/test/v1/service.proto`. + - Pants has source roots at `company/proto/gen` and `company/protogen`, so the + generated Python modules are `services.test.v1.service_pb2` and + `services.test.v1.service_connect`. + """ + rule_runner.set_options( + [ + "--source-root-patterns=['company/proto/gen', 'company/protogen']", + "--python-enable-resolves", + ] + ) + rule_runner.write_files( + { + "buf.yaml": dedent( + """\ + version: v2 + modules: + - path: company/proto + """ + ), + "buf.gen.yaml": dedent( + """\ + version: v2 + managed: + enabled: true + plugins: + - remote: buf.build/protocolbuffers/python:v34.1 + out: company/proto/gen + - remote: buf.build/protocolbuffers/pyi:v34.1 + out: company/proto/gen + - remote: buf.build/connectrpc/python:v0.1.0 + out: company/protogen + """ + ), + "company/proto/services/test/v1/service.proto": "", + "company/proto/services/test/v1/BUILD": ("protobuf_sources(protobuf_generator='buf')"), + } + ) + result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()]) + + address = Address("company/proto/services/test/v1", relative_file_path="service.proto") + provider = (ModuleProvider(address, ModuleProviderType.IMPL),) + assert result == FirstPartyPythonMappingImpl.create( + { + "python-default": { + "services.test.v1.service_pb2": provider, + "services.test.v1.service_connect": provider, + } + } + ) + + +def test_buf_target_grpc_plugin_registers_pb2_grpc(rule_runner: RuleRunner) -> None: + """A grpc-python plugin entry in `buf.gen.yaml` registers `_pb2_grpc` modules.""" + rule_runner.set_options( + [ + "--source-root-patterns=['/']", + "--python-enable-resolves", + ] + ) + rule_runner.write_files( + { + "buf.yaml": "version: v2\nmodules:\n - path: .\n", + "buf.gen.yaml": ( + "version: v2\n" + "plugins:\n" + " - protoc_builtin: python\n" + " out: gen\n" + " - local: protoc-gen-grpc-python\n" + " out: gen\n" + ), + "protos/svc.proto": "", + "protos/BUILD": "protobuf_sources(grpc=True, protobuf_generator='buf')", + } + ) + result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()]) + + def providers(addresses: list[Address]) -> tuple[ModuleProvider, ...]: + return tuple(ModuleProvider(addr, ModuleProviderType.IMPL) for addr in addresses) + + address = Address("protos", relative_file_path="svc.proto") + assert result == FirstPartyPythonMappingImpl.create( + { + "python-default": { + "gen.protos.svc_pb2": providers([address]), + "gen.protos.svc_pb2_grpc": providers([address]), + } + } + ) + + +def test_buf_target_per_target_template_override(rule_runner: RuleRunner) -> None: + """Two buf targets pointing at different `buf_gen_template`s land at different `out:`s.""" + rule_runner.set_options( + [ + "--source-root-patterns=['/']", + "--python-enable-resolves", + ] + ) + rule_runner.write_files( + { + "buf.yaml": "version: v2\nmodules:\n - path: .\n", + "a/protos/svc.proto": "", + "a/protos/buf.gen.yaml": ( + "version: v2\nplugins:\n - protoc_builtin: python\n out: gen_a\n" + ), + "a/protos/BUILD": ( + "protobuf_sources(protobuf_generator='buf', buf_gen_template='buf.gen.yaml')" + ), + "b/protos/svc.proto": "", + "b/protos/buf.gen.yaml": ( + "version: v2\nplugins:\n - protoc_builtin: python\n out: gen_b\n" + ), + "b/protos/BUILD": ( + "protobuf_sources(protobuf_generator='buf', buf_gen_template='buf.gen.yaml')" + ), + } + ) + result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()]) + + def providers(addresses: list[Address]) -> tuple[ModuleProvider, ...]: + return tuple(ModuleProvider(addr, ModuleProviderType.IMPL) for addr in addresses) + + # `buf.yaml` declares one module spanning the repo (`path: .`). Each proto's + # module-relative path is its full repo path, so the generated module is + # `/`. + assert result == FirstPartyPythonMappingImpl.create( + { + "python-default": { + "gen_a.a.protos.svc_pb2": providers( + [Address("a/protos", relative_file_path="svc.proto")] + ), + "gen_b.b.protos.svc_pb2": providers( + [Address("b/protos", relative_file_path="svc.proto")] + ), + } + } + ) + + +def test_buf_target_grpc_field_is_no_op(rule_runner: RuleRunner) -> None: + """`grpc=True` on a buf target has no effect on which suffixes are registered; + plugin presence in `buf.gen.yaml` is the sole determinant.""" + rule_runner.set_options(["--source-root-patterns=['/']", "--python-enable-resolves"]) + rule_runner.write_files( + { + "buf.yaml": "version: v2\nmodules:\n - path: .\n", + # Only `_pb2` plugin — no service plugin. + "buf.gen.yaml": ("version: v2\nplugins:\n - protoc_builtin: python\n out: gen\n"), + "protos/svc.proto": "", + # `grpc=True` would normally register `_pb2_grpc`; on the buf path it's + # ignored and only `_pb2` (matched by the python plugin) is registered. + "protos/BUILD": "protobuf_sources(grpc=True, protobuf_generator='buf')", + } + ) + result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()]) + address = Address("protos", relative_file_path="svc.proto") + provider = (ModuleProvider(address, ModuleProviderType.IMPL),) + assert result == FirstPartyPythonMappingImpl.create( + {"python-default": {"gen.protos.svc_pb2": provider}} + ) + + +def test_buf_target_unpinned_remote_plugin_succeeds_via_registry( + rule_runner: RuleRunner, +) -> None: + """Inference does not enforce pinning — that's codegen's job. An unpinned + `remote:` entry whose plugin id is in Pants's built-in registry is matched + by id, the suffix is found, and the module is registered. (Codegen will + fill in the registry's default `:vX.Y:revN` pin before invoking buf.)""" + rule_runner.set_options(["--source-root-patterns=['/']", "--python-enable-resolves"]) + rule_runner.write_files( + { + "buf.yaml": "version: v2\nmodules:\n - path: .\n", + "buf.gen.yaml": ( + "version: v2\n" + "plugins:\n" + " - remote: buf.build/protocolbuffers/python\n" # no version pin + " out: gen\n" + ), + "protos/svc.proto": "", + "protos/BUILD": "protobuf_sources(protobuf_generator='buf')", + } + ) + result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()]) + address = Address("protos", relative_file_path="svc.proto") + provider = (ModuleProvider(address, ModuleProviderType.IMPL),) + assert result == FirstPartyPythonMappingImpl.create( + {"python-default": {"gen.protos.svc_pb2": provider}} + ) + + +def test_buf_target_bsr_dep_modules_registered_when_include_imports( + rule_runner: RuleRunner, +) -> None: + """When `buf.yaml` has `deps:` pointing at a BSR module Pants knows and + `buf.gen.yaml` sets `include_imports: true` on protocolbuffers/python, the + BSR module's `*_pb2` modules are registered as owned by the proto target.""" + rule_runner.set_options(["--source-root-patterns=['/']", "--python-enable-resolves"]) + rule_runner.write_files( + { + "buf.yaml": ( + "version: v2\nmodules:\n - path: .\ndeps:\n - buf.build/bufbuild/protovalidate\n" + ), + "buf.gen.yaml": ( + "version: v2\n" + "plugins:\n" + " - remote: buf.build/protocolbuffers/python\n" + " out: gen\n" + " include_imports: true\n" + ), + "protos/svc.proto": "", + "protos/BUILD": "protobuf_sources(protobuf_generator='buf')", + } + ) + result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()]) + address = Address("protos", relative_file_path="svc.proto") + provider = (ModuleProvider(address, ModuleProviderType.IMPL),) + expected = { + "gen.protos.svc_pb2": provider, + "buf.validate.expression_pb2": provider, + "buf.validate.validate_pb2": provider, + } + assert result == FirstPartyPythonMappingImpl.create({"python-default": expected}) + + +def test_buf_target_bsr_dep_modules_skipped_without_include_imports( + rule_runner: RuleRunner, +) -> None: + """Without `include_imports: true`, buf doesn't actually generate the BSR + bindings, so we must NOT register them — registering would lie about which + target owns the file and the consumer would still ImportError at runtime.""" + rule_runner.set_options(["--source-root-patterns=['/']", "--python-enable-resolves"]) + rule_runner.write_files( + { + "buf.yaml": ( + "version: v2\nmodules:\n - path: .\ndeps:\n - buf.build/bufbuild/protovalidate\n" + ), + "buf.gen.yaml": ( + "version: v2\n" + "plugins:\n" + " - remote: buf.build/protocolbuffers/python\n" + " out: gen\n" + ), + "protos/svc.proto": "", + "protos/BUILD": "protobuf_sources(protobuf_generator='buf')", + } + ) + result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()]) + address = Address("protos", relative_file_path="svc.proto") + provider = (ModuleProvider(address, ModuleProviderType.IMPL),) + # Only the user's first-party module is registered, not the BSR-dep modules. + assert result == FirstPartyPythonMappingImpl.create( + {"python-default": {"gen.protos.svc_pb2": provider}} + ) + + +def test_buf_target_extra_buf_bsr_modules_extends_registry(rule_runner: RuleRunner) -> None: + """`[python-protobuf].extra_buf_bsr_modules` lets users add BSR module ids + that aren't in Pants's built-in registry.""" + rule_runner.set_options( + [ + "--source-root-patterns=['/']", + "--python-enable-resolves", + ( + "--python-protobuf-extra-buf-bsr-modules=" + '{"buf.build/myorg/internal-types": ' + '["myorg.internal.types.foo_pb2", "myorg.internal.types.bar_pb2"]}' + ), + ] + ) + rule_runner.write_files( + { + "buf.yaml": ( + "version: v2\nmodules:\n - path: .\ndeps:\n - buf.build/myorg/internal-types\n" + ), + "buf.gen.yaml": ( + "version: v2\n" + "plugins:\n" + " - remote: buf.build/protocolbuffers/python\n" + " out: gen\n" + " include_imports: true\n" + ), + "protos/svc.proto": "", + "protos/BUILD": "protobuf_sources(protobuf_generator='buf')", + } + ) + result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()]) + address = Address("protos", relative_file_path="svc.proto") + provider = (ModuleProvider(address, ModuleProviderType.IMPL),) + expected = { + "gen.protos.svc_pb2": provider, + "myorg.internal.types.foo_pb2": provider, + "myorg.internal.types.bar_pb2": provider, + } + assert result == FirstPartyPythonMappingImpl.create({"python-default": expected}) + + +def test_buf_target_extra_plugin_suffixes_override(rule_runner: RuleRunner) -> None: + """`[python-protobuf].extra_buf_plugin_suffixes` lets users teach Pants about + custom or forked plugins without modifying the registry.""" + rule_runner.set_options( + [ + "--source-root-patterns=['/']", + "--python-enable-resolves", + ( + "--python-protobuf-extra-buf-plugin-suffixes=" + '{"remote:myorg.example.com/internal/python-fork": "_pb2"}' + ), + ] + ) + rule_runner.write_files( + { + "buf.yaml": "version: v2\nmodules:\n - path: .\n", + "buf.gen.yaml": ( + "version: v2\n" + "plugins:\n" + " - remote: myorg.example.com/internal/python-fork:v1.0\n" + " out: gen\n" + ), + "protos/svc.proto": "", + "protos/BUILD": "protobuf_sources(protobuf_generator='buf')", + } + ) + result = rule_runner.request(FirstPartyPythonMappingImpl, [PythonProtobufMappingMarker()]) + address = Address("protos", relative_file_path="svc.proto") + provider = (ModuleProvider(address, ModuleProviderType.IMPL),) + assert result == FirstPartyPythonMappingImpl.create( + {"python-default": {"gen.protos.svc_pb2": provider}} + ) + + def test_mypy_protobuf_modules_with_resolves(rule_runner: RuleRunner) -> None: """Verify mypy-protobuf does not change module mapping across resolves.""" rule_runner.set_options( diff --git a/src/python/pants/backend/codegen/protobuf/python/python_protobuf_subsystem.py b/src/python/pants/backend/codegen/protobuf/python/python_protobuf_subsystem.py index 1d4210889c2..0a1c80e2f72 100644 --- a/src/python/pants/backend/codegen/protobuf/python/python_protobuf_subsystem.py +++ b/src/python/pants/backend/codegen/protobuf/python/python_protobuf_subsystem.py @@ -1,11 +1,19 @@ # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md). # Licensed under the Apache License, Version 2.0 (see LICENSE). +from collections.abc import Mapping from dataclasses import dataclass +from pants.backend.codegen.protobuf.buf.config import ( + gen_template_request_from_fields, + parse_plugin_outs, +) +from pants.backend.codegen.protobuf.buf.fields import BufGenTemplateField +from pants.backend.codegen.protobuf.buf.subsystem import BufSubsystem from pants.backend.codegen.protobuf.python.additional_fields import ProtobufPythonResolveField from pants.backend.codegen.protobuf.target_types import ( ProtobufDependenciesField, + ProtobufGeneratorField, ProtobufGrpcToggleField, ) from pants.backend.codegen.utils import find_python_runtime_library_or_raise_error @@ -19,10 +27,13 @@ ) from pants.backend.python.subsystems.python_tool_base import PythonToolRequirementsBase from pants.backend.python.subsystems.setup import PythonSetup +from pants.core.util_rules.config_files import find_config_file +from pants.engine.addresses import Address +from pants.engine.intrinsics import get_digest_contents from pants.engine.rules import collect_rules, implicitly, rule from pants.engine.target import FieldSet, InferDependenciesRequest, InferredDependencies from pants.engine.unions import UnionRule -from pants.option.option_types import BoolOption +from pants.option.option_types import BoolOption, DictOption from pants.option.subsystem import Subsystem from pants.source.source_root import SourceRootRequest, get_source_root from pants.util.docutil import doc_url @@ -32,6 +43,71 @@ # pants: infer-dep(mypy_protobuf.lock*) +# Built-in registry mapping known buf plugin ids to the Python module-name suffix +# their output uses (e.g. `buf.build/grpc/python` produces `*_pb2_grpc.py`, so +# `_pb2_grpc` is the suffix). The Python `python_protobuf_module_mapper` consumes +# this to know which generated modules to register per proto file. Keys are +# `:`, where `` is `remote`, `protoc_builtin`, or `local` — +# matching the field name in `buf.gen.yaml` — so that identical names across +# kinds (e.g. a `local:` plugin named `python`) cannot collide. Users with +# custom plugin ids should layer additional entries via +# `[python-protobuf].extra_buf_plugin_suffixes`. +DEFAULT_PLUGIN_SUFFIXES: Mapping[str, str] = { + # Core message codegen. + "remote:buf.build/protocolbuffers/python": "_pb2", + "protoc_builtin:python": "_pb2", + "local:protoc-gen-python": "_pb2", + # `.pyi` stubs share the same module name as `_pb2.py`, so map to the same suffix. + "remote:buf.build/protocolbuffers/pyi": "_pb2", + "protoc_builtin:pyi": "_pb2", + # ConnectRPC. + "remote:buf.build/connectrpc/python": "_connect", + "local:protoc-gen-connect-python": "_connect", + # gRPC (grpcio). + "remote:buf.build/grpc/python": "_pb2_grpc", + "local:protoc-gen-grpc-python": "_pb2_grpc", + "local:protoc-gen-grpc_python": "_pb2_grpc", + # gRPC (grpclib). + "local:protoc-gen-grpclib_python": "_grpc", +} + + +# Built-in registry mapping known BSR module ids (declared in `buf.yaml`'s +# `deps:`) to the Python module names their generated `*_pb2.py` files produce +# when buf runs with `include_imports: true` on the `_pb2`-emitting plugin. +# The buf module mapper uses this to register dep-inference owners for imports +# of BSR-provided modules from hand-written user code. To refresh: +# `pants export-codegen ::` against a buf.yaml that lists the dep, then list +# the resulting `/<…>_pb2.py` paths. +DEFAULT_BSR_DEP_MODULES: Mapping[str, tuple[str, ...]] = { + "buf.build/bufbuild/protovalidate": ( + "buf.validate.expression_pb2", + "buf.validate.validate_pb2", + ), + "buf.build/protocolbuffers/wellknowntypes": ( + "google.protobuf.any_pb2", + "google.protobuf.api_pb2", + "google.protobuf.descriptor_pb2", + "google.protobuf.duration_pb2", + "google.protobuf.empty_pb2", + "google.protobuf.field_mask_pb2", + "google.protobuf.source_context_pb2", + "google.protobuf.struct_pb2", + "google.protobuf.timestamp_pb2", + "google.protobuf.type_pb2", + "google.protobuf.wrappers_pb2", + ), + "buf.build/googleapis/googleapis": ( + "google.api.annotations_pb2", + "google.api.field_behavior_pb2", + "google.api.http_pb2", + "google.rpc.code_pb2", + "google.rpc.error_details_pb2", + "google.rpc.status_pb2", + ), +} + + class PythonProtobufSubsystem(Subsystem): options_scope = "python-protobuf" help = help_text( @@ -88,6 +164,68 @@ class PythonProtobufSubsystem(Subsystem): ), ) + extra_buf_plugin_suffixes = DictOption[str]( + default={}, + help=softwrap( + """ + Map of additional `buf.gen.yaml` plugin ids to the Python module-name + suffix their output uses, layered on top of Pants's built-in registry + of common plugins (e.g. `buf.build/protocolbuffers/python`, + `buf.build/connectrpc/python`). + + Use this to teach Pants about custom or forked plugins. Keys are + `:`, where `` is `remote`, `protoc_builtin`, or + `local` — matching the field name in the `buf.gen.yaml` plugin entry — + and `` is the plugin id exactly as it appears in that field + (without any `:vX.Y` version suffix on `remote:` entries). Values are + module-name suffixes from the set: + + - `_pb2` — produces message modules (`*_pb2.py`). + - `_pb2_grpc` — produces grpcio service stubs (`*_pb2_grpc.py`). + - `_grpc` — produces grpclib service stubs (`*_grpc.py`). + - `_connect` — produces ConnectRPC service stubs (`*_connect.py`). + + Example: + + extra_buf_plugin_suffixes = { + "remote:myorg.example.com/internal/python-fork": "_pb2", + "remote:buf.build/example/some-grpc-fork": "_pb2_grpc", + "local:protoc-gen-myorg-python": "_pb2", + } + """ + ), + advanced=True, + ) + + extra_buf_bsr_modules = DictOption[list[str]]( + default={}, + help=softwrap( + """ + Map of BSR module ids (declared in `buf.yaml`'s `deps:`) to the + Python module names their generated `*_pb2.py` files produce when + buf runs with `include_imports: true`. Layered on top of Pants's + built-in registry of common modules (e.g. + `buf.build/bufbuild/protovalidate`, + `buf.build/protocolbuffers/wellknowntypes`). + + Use this to teach Pants about BSR deps not in the built-in + registry — typically internal company modules — so that + hand-written code importing them doesn't trip "cannot infer + owners" warnings. + + Example: + + extra_buf_bsr_modules = { + "buf.build/myorg/internal-types": [ + "myorg.internal.types.foo_pb2", + "myorg.internal.types.bar_pb2", + ], + } + """ + ), + advanced=True, + ) + infer_runtime_dependency = BoolOption( default=True, help=softwrap( @@ -138,23 +276,66 @@ class PythonProtobufDependenciesInferenceFieldSet(FieldSet): ProtobufDependenciesField, ProtobufPythonResolveField, ProtobufGrpcToggleField, + ProtobufGeneratorField, + BufGenTemplateField, ) dependencies: ProtobufDependenciesField python_resolve: ProtobufPythonResolveField grpc_toggle: ProtobufGrpcToggleField + generator: ProtobufGeneratorField + buf_gen_template: BufGenTemplateField class InferPythonProtobufDependencies(InferDependenciesRequest): infer_from = PythonProtobufDependenciesInferenceFieldSet +# Mapping from generated-module suffix → (importable module, PyPI requirement +# name, requirement URL). Used by the buf branch of runtime-dep inference to add +# a runtime requirement on the right Python package when a plugin producing that +# suffix appears in `buf.gen.yaml`. +_BUF_RUNTIME_DEPS: tuple[tuple[str, str, str, str], ...] = ( + ("_pb2_grpc", "grpc", "grpcio", "https://pypi.org/project/grpcio/"), + ("_grpc", "grpclib", "grpclib[protobuf]", "https://pypi.org/project/grpclib/"), + ("_connect", "connectrpc", "connectrpc", "https://pypi.org/project/connectrpc/"), +) + + +async def _runtime_dep_for_module( + *, + module: str, + field_set: PythonProtobufDependenciesInferenceFieldSet, + python_setup: PythonSetup, + locality: str | None, + resolve: str, + recommended_requirement_name: str, + recommended_requirement_url: str, + disable_inference_option: str, +) -> Address: + addresses = await map_module_to_address( + PythonModuleOwnersRequest(module, resolve=resolve, locality=locality), + **implicitly(), + ) + return find_python_runtime_library_or_raise_error( + addresses, + field_set.address, + module, + resolve=resolve, + resolves_enabled=python_setup.enable_resolves, + recommended_requirement_name=recommended_requirement_name, + recommended_requirement_url=recommended_requirement_url, + disable_inference_option=disable_inference_option, + ) + + @rule async def infer_dependencies( request: InferPythonProtobufDependencies, python_protobuf: PythonProtobufSubsystem, python_setup: PythonSetup, python_infer_subsystem: PythonInferSubsystem, + buf: BufSubsystem, ) -> InferredDependencies: if not python_protobuf.infer_runtime_dependency: return InferredDependencies([]) @@ -168,74 +349,89 @@ async def infer_dependencies( ) locality = source_root.path + disable_option = f"[{python_protobuf.options_scope}].infer_runtime_dependency" result = [] - addresses_for_protobuf = await map_module_to_address( - PythonModuleOwnersRequest( - "google.protobuf", - resolve=resolve, - locality=locality, - ), - **implicitly(), - ) - result.append( - find_python_runtime_library_or_raise_error( - addresses_for_protobuf, - request.field_set.address, - "google.protobuf", + await _runtime_dep_for_module( + module="google.protobuf", + field_set=request.field_set, + python_setup=python_setup, + locality=locality, resolve=resolve, - resolves_enabled=python_setup.enable_resolves, recommended_requirement_name="protobuf", recommended_requirement_url="https://pypi.org/project/protobuf/", - disable_inference_option=f"[{python_protobuf.options_scope}].infer_runtime_dependency", + disable_inference_option=disable_option, ) ) - if request.field_set.grpc_toggle.value: - if python_protobuf.grpcio_plugin: - addresses_for_grpc = await map_module_to_address( - PythonModuleOwnersRequest( - "grpc", - resolve=resolve, - locality=locality, - ), - **implicitly(), + if request.field_set.generator.value == "buf": + # Buf path: generated-module suffixes come from `buf.gen.yaml`. Subsystem + # booleans and `grpc=True` are not consulted. + template_request = gen_template_request_from_fields( + spec_path=request.field_set.address.spec_path, + address_str=str(request.field_set.address), + override=request.field_set.buf_gen_template.value, + buf=buf, + ) + template_files = await find_config_file(template_request) + suffix_outs: dict[str, str] = {} + if template_files.snapshot.files: + template_path = template_files.snapshot.files[0] + template_dcs = await get_digest_contents(template_files.snapshot.digest) + content = next( + (dc.content for dc in template_dcs if dc.path == template_path), + b"", + ) + # Don't enforce pinning here — codegen does. Inference reads the file + # only to learn plugin ids → suffixes, which works regardless of pin state. + suffix_outs = parse_plugin_outs( + content, + {**DEFAULT_PLUGIN_SUFFIXES, **python_protobuf.extra_buf_plugin_suffixes}, ) + for suffix, module, req_name, req_url in _BUF_RUNTIME_DEPS: + if suffix in suffix_outs: + result.append( + await _runtime_dep_for_module( + module=module, + field_set=request.field_set, + python_setup=python_setup, + locality=locality, + resolve=resolve, + recommended_requirement_name=req_name, + recommended_requirement_url=req_url, + disable_inference_option=disable_option, + ) + ) + return InferredDependencies(result) + # Protoc path: gated on `grpc=True` and the subsystem booleans, since Pants + # drives the protoc invocation directly. + if request.field_set.grpc_toggle.value: + if python_protobuf.grpcio_plugin: result.append( - find_python_runtime_library_or_raise_error( - addresses_for_grpc, - request.field_set.address, + await _runtime_dep_for_module( # Note that the library is called `grpcio`, but the module is `grpc`. - "grpc", + module="grpc", + field_set=request.field_set, + python_setup=python_setup, + locality=locality, resolve=resolve, - resolves_enabled=python_setup.enable_resolves, recommended_requirement_name="grpcio", recommended_requirement_url="https://pypi.org/project/grpcio/", - disable_inference_option=f"[{python_protobuf.options_scope}].infer_runtime_dependency", + disable_inference_option=disable_option, ) ) - if python_protobuf.grpclib_plugin: - addresses_for_grpclib = await map_module_to_address( - PythonModuleOwnersRequest( - "grpclib", - resolve=resolve, - locality=locality, - ), - **implicitly(), - ) - result.append( - find_python_runtime_library_or_raise_error( - addresses_for_grpclib, - request.field_set.address, - "grpclib", + await _runtime_dep_for_module( + module="grpclib", + field_set=request.field_set, + python_setup=python_setup, + locality=locality, resolve=resolve, - resolves_enabled=python_setup.enable_resolves, recommended_requirement_name="grpclib[protobuf]", recommended_requirement_url="https://pypi.org/project/grpclib/", - disable_inference_option=f"[{python_protobuf.options_scope}].infer_runtime_dependency", + disable_inference_option=disable_option, ) ) diff --git a/src/python/pants/backend/codegen/protobuf/python/python_protobuf_subsystem_test.py b/src/python/pants/backend/codegen/protobuf/python/python_protobuf_subsystem_test.py index 4c0f88bb75f..444b8c5879d 100644 --- a/src/python/pants/backend/codegen/protobuf/python/python_protobuf_subsystem_test.py +++ b/src/python/pants/backend/codegen/protobuf/python/python_protobuf_subsystem_test.py @@ -3,6 +3,7 @@ from textwrap import dedent from pants.backend.codegen.protobuf import target_types +from pants.backend.codegen.protobuf.buf import fields as buf_fields from pants.backend.codegen.protobuf.python import additional_fields, python_protobuf_subsystem from pants.backend.codegen.protobuf.python.python_protobuf_subsystem import ( InferPythonProtobufDependencies, @@ -29,6 +30,7 @@ def test_find_protobuf_python_requirement() -> None: *module_mapper.rules(), *stripped_source_files.rules(), *additional_fields.rules(), + *buf_fields.rules(), QueryRule(InferredDependencies, (InferPythonProtobufDependencies,)), ], target_types=[ProtobufSourcesGeneratorTarget, PythonRequirementTarget], @@ -112,6 +114,74 @@ def test_find_protobuf_python_requirement() -> None: ) +def test_buf_target_infers_runtime_deps_from_gen_yaml_plugins() -> None: + """For `protobuf_generator='buf'` targets, runtime-dep inference is driven by + the plugins present in `buf.gen.yaml`, not by `grpc=True`. Connect, grpcio, + and grpclib plugins each contribute their respective runtime requirement.""" + rule_runner = RuleRunner( + rules=[ + *python_protobuf_subsystem.rules(), + *target_types.rules(), + *module_mapper.rules(), + *stripped_source_files.rules(), + *additional_fields.rules(), + *buf_fields.rules(), + QueryRule(InferredDependencies, (InferPythonProtobufDependencies,)), + ], + target_types=[ProtobufSourcesGeneratorTarget, PythonRequirementTarget], + ) + + rule_runner.write_files( + { + "buf.yaml": "version: v2\nmodules:\n - path: .\n", + "buf.gen.yaml": dedent( + """\ + version: v2 + plugins: + - remote: buf.build/protocolbuffers/python:v34.1 + revision: 1 + out: gen + - remote: buf.build/connectrpc/python:v0.10.0 + revision: 1 + out: gen + - remote: buf.build/grpc/python:v1.80.0 + revision: 1 + out: gen + """ + ), + "codegen/dir/f.proto": "", + "codegen/dir/BUILD": "protobuf_sources(protobuf_generator='buf')", + # 3rdparty stand-ins for the runtimes. + "thirdparty/BUILD": dedent( + """\ + python_requirement(name='protobuf', requirements=['protobuf']) + python_requirement(name='grpcio', requirements=['grpcio']) + python_requirement(name='connectrpc', requirements=['connectrpc']) + """ + ), + } + ) + rule_runner.set_options( + [ + "--python-resolves={'python-default': ''}", + "--python-enable-resolves", + "--no-python-enable-lockfile-targets", + ] + ) + + proto_tgt = rule_runner.get_target(Address("codegen/dir", relative_file_path="f.proto")) + request = InferPythonProtobufDependencies( + PythonProtobufDependenciesInferenceFieldSet.create(proto_tgt) + ) + inferred = rule_runner.request(InferredDependencies, [request]) + # protobuf is always inferred; grpcio + connectrpc come from the plugins. + assert set(inferred.include) == { + Address("thirdparty", target_name="protobuf"), + Address("thirdparty", target_name="grpcio"), + Address("thirdparty", target_name="connectrpc"), + } + + def test_find_protobuf_grpclib_python_requirement() -> None: rule_runner = RuleRunner( rules=[ @@ -120,6 +190,7 @@ def test_find_protobuf_grpclib_python_requirement() -> None: *module_mapper.rules(), *stripped_source_files.rules(), *additional_fields.rules(), + *buf_fields.rules(), QueryRule(InferredDependencies, (InferPythonProtobufDependencies,)), ], target_types=[ProtobufSourcesGeneratorTarget, PythonRequirementTarget], diff --git a/src/python/pants/backend/codegen/protobuf/python/register.py b/src/python/pants/backend/codegen/protobuf/python/register.py index 3211453c390..b99f969e483 100644 --- a/src/python/pants/backend/codegen/protobuf/python/register.py +++ b/src/python/pants/backend/codegen/protobuf/python/register.py @@ -8,6 +8,7 @@ from pants.backend.codegen.protobuf import protobuf_dependency_inference from pants.backend.codegen.protobuf import tailor as protobuf_tailor +from pants.backend.codegen.protobuf.buf import fields as buf_fields from pants.backend.codegen.protobuf.python import ( additional_fields, python_protobuf_module_mapper, @@ -25,6 +26,7 @@ def rules(): return [ + *buf_fields.rules(), *additional_fields.rules(), *python_protobuf_subsystem.rules(), *python_rules(), diff --git a/src/python/pants/backend/codegen/protobuf/python/rules.py b/src/python/pants/backend/codegen/protobuf/python/rules.py index 95f9bbdf63f..aed703e92d0 100644 --- a/src/python/pants/backend/codegen/protobuf/python/rules.py +++ b/src/python/pants/backend/codegen/protobuf/python/rules.py @@ -6,14 +6,23 @@ from pants.backend.codegen.protobuf import protoc from pants.backend.codegen.protobuf.protoc import Protoc +from pants.backend.codegen.protobuf.python import buf_rules from pants.backend.codegen.protobuf.python.additional_fields import PythonSourceRootField +from pants.backend.codegen.protobuf.python.buf_rules import ( + GeneratePythonFromProtobufViaBufRequest, + generate_python_from_protobuf_via_buf, +) from pants.backend.codegen.protobuf.python.grpc_python_plugin import GrpcPythonPlugin from pants.backend.codegen.protobuf.python.python_protobuf_subsystem import ( PythonProtobufGrpclibPlugin, PythonProtobufMypyPlugin, PythonProtobufSubsystem, ) -from pants.backend.codegen.protobuf.target_types import ProtobufGrpcToggleField, ProtobufSourceField +from pants.backend.codegen.protobuf.target_types import ( + ProtobufGeneratorField, + ProtobufGrpcToggleField, + ProtobufSourceField, +) from pants.backend.python.target_types import PythonSourceField from pants.backend.python.util_rules import pex from pants.backend.python.util_rules.pex import ( @@ -56,6 +65,12 @@ async def generate_python_from_protobuf( pex_environment: PexEnvironment, platform: Platform, ) -> GeneratedSources: + if request.protocol_target.get(ProtobufGeneratorField).value == "buf": + return await generate_python_from_protobuf_via_buf( + GeneratePythonFromProtobufViaBufRequest(protocol_target=request.protocol_target), + **implicitly(), + ) + download_protoc_request = download_external_tool(protoc.get_request(platform)) output_dir = "_generated_files" @@ -239,6 +254,7 @@ def rules(): return [ *collect_rules(), *pex.rules(), + *buf_rules.rules(), UnionRule(GenerateSourcesRequest, GeneratePythonFromProtobufRequest), *protoc.rules(), UnionRule(ExportableTool, GrpcPythonPlugin), diff --git a/src/python/pants/backend/codegen/protobuf/target_types.py b/src/python/pants/backend/codegen/protobuf/target_types.py index c5734984050..fc6075d8e44 100644 --- a/src/python/pants/backend/codegen/protobuf/target_types.py +++ b/src/python/pants/backend/codegen/protobuf/target_types.py @@ -11,6 +11,7 @@ MultipleSourcesField, OverridesField, SingleSourceField, + StringField, Target, TargetFilesGenerator, TargetFilesGeneratorSettings, @@ -35,6 +36,32 @@ class ProtobufGrpcToggleField(BoolField): help = "Whether to generate gRPC code or not." +class ProtobufGeneratorField(StringField): + alias = "protobuf_generator" + valid_choices = ("protoc", "buf") + default = "protoc" + help = help_text( + """ + Which tool to use to generate code from this `.proto`. Applies to every + language backend that consumes the target. + + - `protoc` (default): use the `protoc` compiler. Output paths follow Pants's + source-root conventions and any per-language overrides + (e.g. `python_source_root`). + - `buf`: use `buf generate` with a `buf.gen.yaml` template. Plugins, output + paths, and managed-mode rewrites come from the template, not from Pants. + The template is resolved per-target via `buf_gen_template`, falling back + to `[buf].gen_template`, falling back to discovery of `buf.gen.yaml` at + the repository root. + + A single `buf.gen.yaml` typically declares plugins for several languages, so + this is a target-level choice rather than a per-language one. To get + protoc-style behavior for an individual language inside a buf-driven build, + declare a `protoc_builtin:` plugin entry in `buf.gen.yaml`. + """ + ) + + class AllProtobufTargets(Targets): pass @@ -60,6 +87,7 @@ class ProtobufSourceTarget(Target): ProtobufDependenciesField, ProtobufSourceField, ProtobufGrpcToggleField, + ProtobufGeneratorField, ) help = help_text( f""" @@ -123,6 +151,7 @@ class ProtobufSourcesGeneratorTarget(TargetFilesGenerator): copied_fields = COMMON_TARGET_FIELDS moved_fields = ( ProtobufGrpcToggleField, + ProtobufGeneratorField, ProtobufDependenciesField, ) settings_request_cls = GeneratorSettingsRequest