From c22f90cf56e1d124fa37e56b49a25986ddc25f2d Mon Sep 17 00:00:00 2001 From: Daniel Muller Date: Thu, 9 May 2024 11:11:06 -0600 Subject: [PATCH] Support protoc-gen-prost-serde in prost_toolchain This allows users to add serde support for their generated prost proto bindings. --- proto/prost/private/prost.bzl | 23 ++++++++++++ proto/prost/private/protoc_wrapper.rs | 54 +++++++++++++++++++++++---- 2 files changed, 69 insertions(+), 8 deletions(-) diff --git a/proto/prost/private/prost.bzl b/proto/prost/private/prost.bzl index b7d55c362d..fa49fb80d9 100644 --- a/proto/prost/private/prost.bzl +++ b/proto/prost/private/prost.bzl @@ -73,6 +73,14 @@ def _compile_proto(ctx, crate_name, proto_info, deps, prost_toolchain, rustfmt_t additional_args.add_all(prost_toolchain.tonic_opts, format_each = "--tonic_opt=%s") tools = depset([tonic_plugin.executable], transitive = [tools]) + if prost_toolchain.prost_serde_plugin: + prost_serde_plugin = prost_toolchain.prost_serde_plugin[DefaultInfo].files_to_run + additional_args.add(prost_toolchain.prost_serde_plugin_flag % prost_serde_plugin.executable.path) + additional_args.add("--prost-serde_opt=no_include") + additional_args.add("--is_prost_serde") + additional_args.add_all(prost_toolchain.prost_serde_opts, format_each = "--prost-serde_opt=%s") + tools = depset([prost_serde_plugin.executable], transitive = [tools]) + if rustfmt_toolchain: additional_args.add("--rustfmt={}".format(rustfmt_toolchain.rustfmt.path)) tools = depset(transitive = [tools, rustfmt_toolchain.all_files]) @@ -359,6 +367,9 @@ def _rust_prost_toolchain_impl(ctx): tonic_plugin = ctx.attr.tonic_plugin, tonic_plugin_flag = ctx.attr.tonic_plugin_flag, tonic_runtime = ctx.attr.tonic_runtime, + prost_serde_opts = ctx.attr.prost_serde_opts, + prost_serde_plugin = ctx.attr.prost_serde_plugin, + prost_serde_plugin_flag = ctx.attr.prost_serde_plugin_flag, )] rust_prost_toolchain = rule( @@ -411,6 +422,18 @@ rust_prost_toolchain = rule( doc = "The Tonic runtime crates to use.", providers = [[rust_common.crate_info], [rust_common.crate_group_info]], ), + "prost_serde_opts": attr.string_list( + doc = "Additional options to add to prost-serde.", + ), + "prost_serde_plugin": attr.label( + doc = "Additional plugins to add to prost-serde.", + cfg = "exec", + executable = True, + ), + "prost_serde_plugin_flag": attr.string( + doc = "prost-serde plugin flag format. (e.g. `--plugin=protoc-gen-prost-serde=%s`))", + default = "--plugin=protoc-gen-prost-serde=%s", + ), }, ) diff --git a/proto/prost/private/protoc_wrapper.rs b/proto/prost/private/protoc_wrapper.rs index 9c41892ccf..dfd73bf018 100644 --- a/proto/prost/private/protoc_wrapper.rs +++ b/proto/prost/private/protoc_wrapper.rs @@ -117,7 +117,11 @@ struct Module { /// } /// } /// ``` -fn generate_lib_rs(prost_outputs: &BTreeSet, is_tonic: bool) -> String { +fn generate_lib_rs( + prost_outputs: &BTreeSet, + is_tonic: bool, + is_prost_serde: bool, +) -> String { let mut module_info = BTreeMap::new(); for path in prost_outputs.iter() { @@ -133,7 +137,13 @@ fn generate_lib_rs(prost_outputs: &BTreeSet, is_tonic: bool) -> String .strip_suffix(".tonic") .expect("Failed to strip suffix") .to_string() - }; + } + if is_prost_serde && package.ends_with(".serde") { + package = package + .strip_suffix(".serde") + .expect("Failed to strip suffix") + .to_string() + } if package.is_empty() { continue; @@ -155,14 +165,21 @@ fn generate_lib_rs(prost_outputs: &BTreeSet, is_tonic: bool) -> String // Avoid a stack overflow by skipping a known bad package name let module_name = snake_cased_package_name(&package); - module_info.insert( - module_name.clone(), - Module { + module_info + .entry(module_name.clone()) + // Serde module will have same name as module its bindings are for. + .and_modify(|module: &mut Module| { + module.contents = format!( + "{}\n{}", + module.contents, + fs::read_to_string(path).expect("Failed to read file") + ); + }) + .or_insert(Module { name, contents: fs::read_to_string(path).expect("Failed to read file"), submodules: BTreeSet::new(), - }, - ); + }); let module_parts = module_name.split('.').collect::>(); for parent_module_index in 0..module_parts.len() { @@ -443,6 +460,9 @@ struct Args { /// Whether to generate tonic code. is_tonic: bool, + /// Whether to generate prost serde code. + is_prost_serde: bool, + /// Extra arguments to pass to protoc. extra_args: Vec, } @@ -463,6 +483,7 @@ impl Args { let mut label: Option = None; let mut tonic_or_prost_opts = Vec::new(); let mut is_tonic = false; + let mut is_prost_serde = false; let mut extra_args = Vec::new(); @@ -486,6 +507,11 @@ impl Args { return; } + if arg == "--is_prost_serde" { + is_prost_serde = true; + return; + } + if !arg.contains('=') { extra_args.push(arg); return; @@ -563,6 +589,9 @@ impl Args { if is_tonic { extra_args.push(format!("--tonic_opt={}", tonic_or_prost_opt)); } + if is_prost_serde { + extra_args.push(format!("--prost-serde_opt={}", tonic_or_prost_opt)); + } } if protoc.is_none() { @@ -613,6 +642,7 @@ impl Args { rustfmt, proto_paths, is_tonic, + is_prost_serde, label: label.unwrap(), extra_args, }) @@ -718,6 +748,7 @@ fn main() { rustfmt, proto_paths, is_tonic, + is_prost_serde, extra_args, } = Args::parse().expect("Failed to parse args"); @@ -737,6 +768,9 @@ fn main() { if is_tonic { cmd.arg(format!("--tonic_out={}", out_dir.display())); } + if is_prost_serde { + cmd.arg(format!("--prost-serde_out={}", out_dir.display())); + } cmd.args(extra_args); cmd.args( proto_paths @@ -830,7 +864,11 @@ fn main() { .expect("Failed to compute proto package info"); // Write outputs - fs::write(&out_librs, generate_lib_rs(&rust_files, is_tonic)).expect("Failed to write file."); + fs::write( + &out_librs, + generate_lib_rs(&rust_files, is_tonic, is_prost_serde), + ) + .expect("Failed to write file."); fs::write( package_info_file, extern_paths