From bab497f38404cd8d7af5ab3ac636a3c203f9df0e Mon Sep 17 00:00:00 2001 From: Lucio Franco Date: Sat, 5 Oct 2019 13:01:01 -0400 Subject: [PATCH] fix(codegen): Use wellknown types from `prost-types` This change introduces the ability to let prost-build provide the fully qualifed type path for types that start with `google.protobuf`. --- tonic-build/src/client.rs | 12 +++------ tonic-build/src/lib.rs | 27 +++++++++++++------- tonic-build/src/server.rs | 15 ++++------- tonic-build/tests/protos/empty.proto | 7 ----- tonic-build/tests/protos/wellknown.proto | 13 ++++++++++ tonic-build/tests/{empty.rs => wellknown.rs} | 4 +-- 6 files changed, 42 insertions(+), 36 deletions(-) delete mode 100644 tonic-build/tests/protos/empty.proto create mode 100644 tonic-build/tests/protos/wellknown.proto rename tonic-build/tests/{empty.rs => wellknown.rs} (61%) diff --git a/tonic-build/src/client.rs b/tonic-build/src/client.rs index 63ef3254d..3056a9585 100644 --- a/tonic-build/src/client.rs +++ b/tonic-build/src/client.rs @@ -96,8 +96,7 @@ fn generate_methods(service: &Service, proto: &str) -> TokenStream { fn generate_unary(method: &Method, proto: &str, path: String) -> TokenStream { let ident = format_ident!("{}", method.name); - let request = crate::replace_wellknown(proto, &method.input_type); - let response = crate::replace_wellknown(proto, &method.output_type); + let (request, response) = crate::replace_wellknown(proto, &method); quote! { pub async fn #ident(&mut self, request: tonic::Request<#request>) @@ -113,8 +112,7 @@ fn generate_unary(method: &Method, proto: &str, path: String) -> TokenStream { fn generate_server_streaming(method: &Method, proto: &str, path: String) -> TokenStream { let ident = format_ident!("{}", method.name); - let request = crate::replace_wellknown(proto, &method.input_type); - let response = crate::replace_wellknown(proto, &method.output_type); + let (request, response) = crate::replace_wellknown(proto, &method); quote! { pub async fn #ident(&mut self, request: tonic::Request<#request>) @@ -130,8 +128,7 @@ fn generate_server_streaming(method: &Method, proto: &str, path: String) -> Toke fn generate_client_streaming(method: &Method, proto: &str, path: String) -> TokenStream { let ident = format_ident!("{}", method.name); - let request = crate::replace_wellknown(proto, &method.input_type); - let response = crate::replace_wellknown(proto, &method.output_type); + let (request, response) = crate::replace_wellknown(proto, &method); quote! { pub async fn #ident(&mut self, request: tonic::Request) @@ -149,8 +146,7 @@ fn generate_client_streaming(method: &Method, proto: &str, path: String) -> Toke fn generate_streaming(method: &Method, proto: &str, path: String) -> TokenStream { let ident = format_ident!("{}", method.name); - let request = crate::replace_wellknown(proto, &method.input_type); - let response = crate::replace_wellknown(proto, &method.output_type); + let (request, response) = crate::replace_wellknown(proto, &method); quote! { pub async fn #ident(&mut self, request: tonic::Request) diff --git a/tonic-build/src/lib.rs b/tonic-build/src/lib.rs index f0d1b1e5d..57f8ff9b0 100644 --- a/tonic-build/src/lib.rs +++ b/tonic-build/src/lib.rs @@ -58,7 +58,7 @@ #![doc(test(no_crate_inject, attr(deny(rust_2018_idioms))))] use proc_macro2::{Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenStream}; -use prost_build::Config; +use prost_build::{Config, Method}; use quote::{ToTokens, TokenStreamExt}; #[cfg(feature = "rustfmt")] @@ -281,13 +281,22 @@ fn generate_doc_comments>(comments: &[T]) -> TokenStream { stream } -fn replace_wellknown(proto_path: &str, output: &str) -> TokenStream { - // TODO: detect more wellknown protobuf types - // https://github.com/danburkert/prost/blob/master/prost-types/src/protobuf.rs - match output { - "()" => quote::quote! { () }, - _ => syn::parse_str::(&format!("{}::{}", proto_path, output)) +fn replace_wellknown(proto_path: &str, method: &Method) -> (TokenStream, TokenStream) { + let request = if method.input_proto_type.starts_with(".google.protobuf") { + method.input_type.parse::().unwrap() + } else { + syn::parse_str::(&format!("{}::{}", proto_path, method.input_type)) .unwrap() - .to_token_stream(), - } + .to_token_stream() + }; + + let response = if method.output_proto_type.starts_with(".google.protobuf") { + method.input_type.parse::().unwrap() + } else { + syn::parse_str::(&format!("{}::{}", proto_path, method.output_type)) + .unwrap() + .to_token_stream() + }; + + (request, response) } diff --git a/tonic-build/src/server.rs b/tonic-build/src/server.rs index bc11cc2e1..d19f45a7d 100644 --- a/tonic-build/src/server.rs +++ b/tonic-build/src/server.rs @@ -114,8 +114,7 @@ fn generate_trait_methods(service: &Service, proto_path: &str) -> TokenStream { for method in &service.methods { let name = quote::format_ident!("{}", method.name); - let req_message = crate::replace_wellknown(proto_path, &method.input_type); - let res_message = crate::replace_wellknown(proto_path, &method.output_type); + let (req_message, res_message) = crate::replace_wellknown(proto_path, &method); let method_doc = generate_doc_comments(&method.comments.leading); @@ -226,8 +225,7 @@ fn generate_unary( ) -> TokenStream { let service_ident = Ident::new(&method.proto_name, Span::call_site()); - let request = crate::replace_wellknown(proto_path, &method.input_type); - let response = crate::replace_wellknown(proto_path, &method.output_type); + let (request, response) = crate::replace_wellknown(proto_path, &method); quote! { struct #service_ident(pub Arc); @@ -266,8 +264,7 @@ fn generate_server_streaming( ) -> TokenStream { let service_ident = Ident::new(&method.proto_name, Span::call_site()); - let request = crate::replace_wellknown(proto_path, &method.input_type); - let response = crate::replace_wellknown(proto_path, &method.output_type); + let (request, response) = crate::replace_wellknown(proto_path, &method); let response_stream = quote::format_ident!("{}Stream", method.proto_name); @@ -310,8 +307,7 @@ fn generate_client_streaming( ) -> TokenStream { let service_ident = Ident::new(&method.proto_name, Span::call_site()); - let request = crate::replace_wellknown(proto_path, &method.input_type); - let response = crate::replace_wellknown(proto_path, &method.output_type); + let (request, response) = crate::replace_wellknown(proto_path, &method); quote! { struct #service_ident(pub Arc); @@ -352,8 +348,7 @@ fn generate_streaming( ) -> TokenStream { let service_ident = Ident::new(&method.proto_name, Span::call_site()); - let request = crate::replace_wellknown(proto_path, &method.input_type); - let response = crate::replace_wellknown(proto_path, &method.output_type); + let (request, response) = crate::replace_wellknown(proto_path, &method); let response_stream = quote::format_ident!("{}Stream", method.proto_name); diff --git a/tonic-build/tests/protos/empty.proto b/tonic-build/tests/protos/empty.proto deleted file mode 100644 index 0d5532ce0..000000000 --- a/tonic-build/tests/protos/empty.proto +++ /dev/null @@ -1,7 +0,0 @@ -syntax = "proto3"; -package empty; -import "google/protobuf/empty.proto"; - -service Admin { - rpc EmptyCall(google.protobuf.Empty) returns (google.protobuf.Empty); -} diff --git a/tonic-build/tests/protos/wellknown.proto b/tonic-build/tests/protos/wellknown.proto new file mode 100644 index 000000000..8ba4093b1 --- /dev/null +++ b/tonic-build/tests/protos/wellknown.proto @@ -0,0 +1,13 @@ +syntax = "proto3"; + +package wellknown; + +import "google/protobuf/empty.proto"; +import "google/protobuf/wrappers.proto"; +import "google/protobuf/any.proto"; + +service Admin { + rpc EmptyCall(google.protobuf.Empty) returns (google.protobuf.Empty); + rpc StringCall(google.protobuf.StringValue) returns (google.protobuf.Empty); + rpc AnyCall(google.protobuf.Any) returns (google.protobuf.Empty); +} diff --git a/tonic-build/tests/empty.rs b/tonic-build/tests/wellknown.rs similarity index 61% rename from tonic-build/tests/empty.rs rename to tonic-build/tests/wellknown.rs index 5e644cd2c..921d7f51d 100644 --- a/tonic-build/tests/empty.rs +++ b/tonic-build/tests/wellknown.rs @@ -1,9 +1,9 @@ #[test] -fn empty() { +fn wellknown() { let tmp = std::env::temp_dir(); tonic_build::configure() .out_dir(tmp) .format(false) - .compile(&["tests/protos/empty.proto"], &["tests/protos"]) + .compile(&["tests/protos/wellknown.proto"], &["tests/protos"]) .unwrap(); }