diff --git a/extensions/prost/private/prost.bzl b/extensions/prost/private/prost.bzl index f06719db95..ec3b029b05 100644 --- a/extensions/prost/private/prost.bzl +++ b/extensions/prost/private/prost.bzl @@ -86,6 +86,9 @@ def _compile_proto( additional_args.add("--additional_srcs={}".format(",".join([f.path for f in all_additional_srcs.to_list()]))) additional_args.add_all(prost_toolchain.prost_opts + prost_opts, format_each = "--prost_opt=%s") + if prost_toolchain.is_no_std: + additional_args.add("--is_no_std") + if prost_toolchain.tonic_plugin: tonic_plugin = prost_toolchain.tonic_plugin[DefaultInfo].files_to_run additional_args.add(prost_toolchain.tonic_plugin_flag % tonic_plugin.executable.path) @@ -442,6 +445,7 @@ def _rust_prost_toolchain_impl(ctx): tonic_plugin_flag = ctx.attr.tonic_plugin_flag, tonic_runtime = ctx.attr.tonic_runtime, include_transitive_deps = ctx.attr.include_transitive_deps, + is_no_std = ctx.attr.is_no_std, )] rust_prost_toolchain = rule( @@ -453,6 +457,10 @@ rust_prost_toolchain = rule( doc = "Whether to include transitive dependencies. If set to True, all transitive dependencies will directly accessible by the dependent crate.", default = False, ), + "is_no_std": attr.bool( + doc = "If a no_std tag should be put into the generated code.", + default = False, + ), "prost_opts": attr.string_list( doc = "Additional options to add to Prost.", ), diff --git a/extensions/prost/private/protoc_wrapper.rs b/extensions/prost/private/protoc_wrapper.rs index a0403f82e8..ded596ac19 100644 --- a/extensions/prost/private/protoc_wrapper.rs +++ b/extensions/prost/private/protoc_wrapper.rs @@ -156,8 +156,17 @@ fn generate_lib_rs( is_tonic: bool, direct_dep_crate_names: Vec, additional_content: String, + is_no_std: bool, ) -> String { - let mut contents = vec!["// @generated".to_string(), "".to_string()]; + let mut contents = vec![ + if is_no_std { + "#![no_std]".to_string() + } else { + "".to_string() + }, + "// @generated".to_string(), + "".to_string(), + ]; for crate_name in direct_dep_crate_names { contents.push(format!("pub use {crate_name};")); } @@ -457,6 +466,9 @@ struct Args { /// Whether to generate tonic code. is_tonic: bool, + // Whether to put a no_std tag into the generated code. + is_no_std: bool, + /// Extra arguments to pass to protoc. extra_args: Vec, } @@ -479,6 +491,7 @@ impl Args { let mut tonic_or_prost_opts = Vec::new(); let mut direct_dep_crate_names = Vec::new(); let mut is_tonic = false; + let mut is_no_std = false; let mut extra_args = Vec::new(); @@ -501,6 +514,10 @@ impl Args { is_tonic = true; return; } + if arg == "--is_no_std" { + is_no_std = true; + return; + } if !arg.contains('=') { extra_args.push(arg); @@ -644,6 +661,7 @@ impl Args { proto_paths, direct_dep_crate_names, is_tonic, + is_no_std, label: label.unwrap(), extra_args, }) @@ -748,6 +766,7 @@ fn main() { proto_paths, direct_dep_crate_names, is_tonic, + is_no_std, extra_args, } = Args::parse().expect("Failed to parse args"); @@ -917,6 +936,7 @@ fn main() { is_tonic, direct_dep_crate_names, additional_content, + is_no_std, ), ) .expect("Failed to write file."); @@ -972,7 +992,6 @@ fn escape_keyword(s: String) -> String { #[cfg(test)] mod test { - use super::*; use prost_types::{FieldDescriptorProto, ServiceDescriptorProto};