diff --git a/examples/build.rs b/examples/build.rs index 9674a9170..8e4f64ff1 100644 --- a/examples/build.rs +++ b/examples/build.rs @@ -15,6 +15,14 @@ fn main() { tonic_build::compile_protos("proto/echo/echo.proto").unwrap(); + tonic_build::configure() + .server_mod_attribute("attrs", "#[cfg(feature = \"server\")]") + .server_attribute("Echo", "#[derive(PartialEq)]") + .client_mod_attribute("attrs", "#[cfg(feature = \"client\")]") + .client_attribute("Echo", "#[derive(PartialEq)]") + .compile(&["proto/attrs/attrs.proto"], &["proto"]) + .unwrap(); + tonic_build::configure() .build_server(false) .compile( diff --git a/examples/proto/attrs/attrs.proto b/examples/proto/attrs/attrs.proto new file mode 100644 index 000000000..b40c8860b --- /dev/null +++ b/examples/proto/attrs/attrs.proto @@ -0,0 +1,19 @@ +syntax = "proto3"; + +package attrs; + +// EchoRequest is the request for echo. +message EchoRequest { + string message = 1; +} + +// EchoResponse is the response for echo. +message EchoResponse { + string message = 1; +} + +// Echo is the echo service. +service Echo { + // UnaryEcho is unary echo. + rpc UnaryEcho(EchoRequest) returns (EchoResponse) {} +} \ No newline at end of file diff --git a/tonic-build/src/client.rs b/tonic-build/src/client.rs index b44c1d0a2..8125c5cd9 100644 --- a/tonic-build/src/client.rs +++ b/tonic-build/src/client.rs @@ -1,4 +1,4 @@ -use super::{Method, Service}; +use super::{Attributes, Method, Service}; use crate::{generate_doc_comments, naive_snake_case}; use proc_macro2::TokenStream; use quote::{format_ident, quote}; @@ -12,6 +12,7 @@ pub fn generate( emit_package: bool, proto_path: &str, compile_well_known_types: bool, + attributes: &Attributes, ) -> TokenStream { let service_ident = quote::format_ident!("{}Client", service.name()); let client_mod = quote::format_ident!("{}_client", naive_snake_case(&service.name())); @@ -20,13 +21,26 @@ pub fn generate( let connect = generate_connect(&service_ident); let service_doc = generate_doc_comments(service.comment()); + let package = if emit_package { service.package() } else { "" }; + let path = format!( + "{}{}{}", + package, + if package.is_empty() { "" } else { "." }, + service.identifier() + ); + + let mod_attributes = attributes.for_mod(package); + let struct_attributes = attributes.for_struct(&path); + quote! { /// Generated client implementations. + #(#mod_attributes)* pub mod #client_mod { #![allow(unused_variables, dead_code, missing_docs)] use tonic::codegen::*; #service_doc + #(#struct_attributes)* #[derive(Debug, Clone)] pub struct #service_ident { inner: tonic::client::Grpc, diff --git a/tonic-build/src/lib.rs b/tonic-build/src/lib.rs index 7d6d67eae..b1d315d65 100644 --- a/tonic-build/src/lib.rs +++ b/tonic-build/src/lib.rs @@ -155,6 +155,53 @@ pub trait Method { ) -> (TokenStream, TokenStream); } +/// Attributes that will be added to `mod` and `struct` items. Should be pairs of (`pattern`, `attribute`). +#[non_exhaustive] +#[derive(Debug, Default, Clone)] +pub struct Attributes { + /// `mod` attributes. + pub module: Vec<(String, String)>, + /// `struct` attributes. + pub structure: Vec<(String, String)>, +} + +impl Attributes { + fn for_mod(&self, name: &str) -> Vec { + generate_attributes(name, &self.module) + } + + fn for_struct(&self, name: &str) -> Vec { + generate_attributes(name, &self.structure) + } + + #[cfg(feature = "prost")] + fn push_mod(&mut self, attr: (String, String)) { + self.module.push(attr); + } + + #[cfg(feature = "prost")] + fn push_struct(&mut self, attr: (String, String)) { + self.structure.push(attr); + } +} + +// Generates attributes given a list of (`pattern`, `attribute`) pairs. If `pattern` matches `name`, `attribute` will be included. +fn generate_attributes<'a>( + name: &str, + attrs: impl IntoIterator, +) -> Vec { + attrs + .into_iter() + .filter(|(matcher, _)| match_name(matcher, name)) + .flat_map(|(_, attr)| { + // attributes cannot be parsed directly, so we pretend they're on a struct + syn::parse_str::(&format!("{}\nstruct fake;", attr)) + .unwrap() + .attrs + }) + .collect::>() +} + /// Format files under the out_dir with rustfmt #[cfg(feature = "rustfmt")] #[cfg_attr(docsrs, doc(cfg(feature = "rustfmt")))] @@ -218,6 +265,34 @@ fn generate_doc_comments>(comments: &[T]) -> TokenStream { stream } +// Checks whether a path pattern matches a given path. +pub(crate) fn match_name(pattern: &str, path: &str) -> bool { + if pattern.is_empty() { + false + } else if pattern == "." { + true + } else if pattern == path { + true + } else { + let pattern_segments = pattern.split('.').collect::>(); + let path_segments = path.split('.').collect::>(); + + if &pattern[..1] == "." { + // prefix match + if pattern_segments.len() > path_segments.len() { + false + } else { + pattern_segments[..] == path_segments[..pattern_segments.len()] + } + // suffix match + } else if pattern_segments.len() > path_segments.len() { + false + } else { + pattern_segments[..] == path_segments[path_segments.len() - pattern_segments.len()..] + } + } +} + fn naive_snake_case(name: &str) -> String { let mut s = String::new(); let mut it = name.chars().peekable(); @@ -234,14 +309,40 @@ fn naive_snake_case(name: &str) -> String { s } -#[test] -fn test_snake_case() { - for case in &[ - ("Service", "service"), - ("ThatHasALongName", "that_has_a_long_name"), - ("greeter", "greeter"), - ("ABCServiceX", "a_b_c_service_x"), - ] { - assert_eq!(naive_snake_case(case.0), case.1) +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_match_name() { + assert!(match_name(".", ".my.protos")); + assert!(match_name(".", ".protos")); + + assert!(match_name(".my", ".my")); + assert!(match_name(".my", ".my.protos")); + assert!(match_name(".my.protos.Service", ".my.protos.Service")); + + assert!(match_name("Service", ".my.protos.Service")); + + assert!(!match_name(".m", ".my.protos")); + assert!(!match_name(".p", ".protos")); + + assert!(!match_name(".my", ".myy")); + assert!(!match_name(".protos", ".my.protos")); + assert!(!match_name(".Service", ".my.protos.Service")); + + assert!(!match_name("service", ".my.protos.Service")); + } + + #[test] + fn test_snake_case() { + for case in &[ + ("Service", "service"), + ("ThatHasALongName", "that_has_a_long_name"), + ("greeter", "greeter"), + ("ABCServiceX", "a_b_c_service_x"), + ] { + assert_eq!(naive_snake_case(case.0), case.1) + } } } diff --git a/tonic-build/src/prost.rs b/tonic-build/src/prost.rs index ee62603d6..e3068b3b0 100644 --- a/tonic-build/src/prost.rs +++ b/tonic-build/src/prost.rs @@ -1,4 +1,4 @@ -use super::{client, server}; +use super::{client, server, Attributes}; use proc_macro2::TokenStream; use prost_build::{Config, Method, Service}; use quote::ToTokens; @@ -18,6 +18,8 @@ pub fn configure() -> Builder { extern_path: Vec::new(), field_attributes: Vec::new(), type_attributes: Vec::new(), + server_attributes: Attributes::default(), + client_attributes: Attributes::default(), proto_path: "super".to_string(), compile_well_known_types: false, #[cfg(feature = "rustfmt")] @@ -162,6 +164,7 @@ impl prost_build::ServiceGenerator for ServiceGenerator { self.builder.emit_package, &self.builder.proto_path, self.builder.compile_well_known_types, + &self.builder.server_attributes, ); self.servers.extend(server); } @@ -172,6 +175,7 @@ impl prost_build::ServiceGenerator for ServiceGenerator { self.builder.emit_package, &self.builder.proto_path, self.builder.compile_well_known_types, + &self.builder.client_attributes, ); self.clients.extend(client); } @@ -215,6 +219,8 @@ pub struct Builder { pub(crate) extern_path: Vec<(String, String)>, pub(crate) field_attributes: Vec<(String, String)>, pub(crate) type_attributes: Vec<(String, String)>, + pub(crate) server_attributes: Attributes, + pub(crate) client_attributes: Attributes, pub(crate) proto_path: String, pub(crate) emit_package: bool, pub(crate) compile_well_known_types: bool, @@ -291,6 +297,42 @@ impl Builder { self } + /// Add additional attribute to matched server `mod`s. Matches on the package name. + pub fn server_mod_attribute, A: AsRef>( + mut self, + path: P, + attribute: A, + ) -> Self { + self.server_attributes + .push_mod((path.as_ref().to_string(), attribute.as_ref().to_string())); + self + } + + /// Add additional attribute to matched service servers. Matches on the service name. + pub fn server_attribute, A: AsRef>(mut self, path: P, attribute: A) -> Self { + self.server_attributes + .push_struct((path.as_ref().to_string(), attribute.as_ref().to_string())); + self + } + + /// Add additional attribute to matched client `mod`s. Matches on the package name. + pub fn client_mod_attribute, A: AsRef>( + mut self, + path: P, + attribute: A, + ) -> Self { + self.client_attributes + .push_mod((path.as_ref().to_string(), attribute.as_ref().to_string())); + self + } + + /// Add additional attribute to matched service clients. Matches on the service name. + pub fn client_attribute, A: AsRef>(mut self, path: P, attribute: A) -> Self { + self.client_attributes + .push_struct((path.as_ref().to_string(), attribute.as_ref().to_string())); + self + } + /// Set the path to where tonic will search for the Request/Response proto structs /// live relative to the module where you call `include_proto!`. /// diff --git a/tonic-build/src/server.rs b/tonic-build/src/server.rs index f8962432f..be10b8302 100644 --- a/tonic-build/src/server.rs +++ b/tonic-build/src/server.rs @@ -1,4 +1,4 @@ -use super::{Method, Service}; +use super::{Attributes, Method, Service}; use crate::{generate_doc_comment, generate_doc_comments, naive_snake_case}; use proc_macro2::{Span, TokenStream}; use quote::quote; @@ -13,6 +13,7 @@ pub fn generate( emit_package: bool, proto_path: &str, compile_well_known_types: bool, + attributes: &Attributes, ) -> TokenStream { let methods = generate_methods(service, proto_path, compile_well_known_types); @@ -35,6 +36,8 @@ pub fn generate( service.identifier() ); let transport = generate_transport(&server_service, &server_trait, &path); + let mod_attributes = attributes.for_mod(package); + let struct_attributes = attributes.for_struct(&path); let compression_enabled = cfg!(feature = "compression"); @@ -64,6 +67,7 @@ pub fn generate( quote! { /// Generated server implementations. + #(#mod_attributes)* pub mod #server_mod { #![allow(unused_variables, dead_code, missing_docs)] use tonic::codegen::*; @@ -71,6 +75,7 @@ pub fn generate( #generated_trait #service_doc + #(#struct_attributes)* #[derive(Debug)] pub struct #server_service { inner: _Inner,