feat(build): support adding attributes to clients and servers (#684)

2021-07-08 15:01:28 -04:00
parent 89df3380dd
commit a948a8f884
6 changed files with 216 additions and 12 deletions
@@ -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(
@@ -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) {}
+}
@@ -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<T: Service>(
    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<T: Service>(
    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<T> {
                inner: tonic::client::Grpc<T>,
@@ -155,6 +155,68 @@ pub trait Method {
    ) -> (TokenStream, TokenStream);
 }

+/// Attributes that will be added to `mod` and `struct` items.
+#[derive(Debug, Default, Clone)]
+pub struct Attributes {
+    /// `mod` attributes.
+    module: Vec<(String, String)>,
+    /// `struct` attributes.
+    structure: Vec<(String, String)>,
+}
+
+impl Attributes {
+    fn for_mod(&self, name: &str) -> Vec<syn::Attribute> {
+        generate_attributes(name, &self.module)
+    }
+
+    fn for_struct(&self, name: &str) -> Vec<syn::Attribute> {
+        generate_attributes(name, &self.structure)
+    }
+
+    /// Add an attribute that will be added to `mod` items matching the given pattern.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use tonic_build::*;
+    /// let mut attributes = Attributes::default();
+    /// attributes.push_mod("my.proto.package", r#"#[cfg(feature = "server")]"#);
+    /// ```
+    pub fn push_mod(&mut self, pattern: impl Into<String>, attr: impl Into<String>) {
+        self.module.push((pattern.into(), attr.into()));
+    }
+
+    /// Add an attribute that will be added to `struct` items matching the given pattern.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use tonic_build::*;
+    /// let mut attributes = Attributes::default();
+    /// attributes.push_struct("EchoService", "#[derive(PartialEq)]");
+    /// ```
+    pub fn push_struct(&mut self, pattern: impl Into<String>, attr: impl Into<String>) {
+        self.structure.push((pattern.into(), attr.into()));
+    }
+}
+
+// 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<Item = &'a (String, String)>,
+) -> Vec<syn::Attribute> {
+    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::<syn::DeriveInput>(&format!("{}\nstruct fake;", attr))
+                .unwrap()
+                .attrs
+        })
+        .collect::<Vec<_>>()
+}
+
 /// Format files under the out_dir with rustfmt
 #[cfg(feature = "rustfmt")]
 #[cfg_attr(docsrs, doc(cfg(feature = "rustfmt")))]
@@ -218,6 +280,34 @@ fn generate_doc_comments<T: AsRef<str>>(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::<Vec<_>>();
+        let path_segments = path.split('.').collect::<Vec<_>>();
+
+        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 +324,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)
+        }
    }
 }
@@ -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<P: AsRef<str>, A: AsRef<str>>(
+        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<P: AsRef<str>, A: AsRef<str>>(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<P: AsRef<str>, A: AsRef<str>>(
+        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<P: AsRef<str>, A: AsRef<str>>(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!`.
    ///
@@ -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<T: Service>(
    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<T: Service>(
        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<T: Service>(

    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<T: Service>(
            #generated_trait

            #service_doc
+            #(#struct_attributes)*
            #[derive(Debug)]
            pub struct #server_service<T: #server_trait> {
                inner: _Inner<T>,