fix(codegen): Use wellknown types from prost-types (#49)

This change introduces the ability to let prost-build
provide the fully qualifed type path for types that
start with `google.protobuf`.

Author: LucioFranco

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<S>(&mut self, request: tonic::Request<S>)
@@ -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<S>(&mut self, request: tonic::Request<S>)

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<T: AsRef<str>>(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::<syn::Path>(&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::<TokenStream>().unwrap()
+    } else {
+        syn::parse_str::<syn::Path>(&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::<TokenStream>().unwrap()
+    } else {
+        syn::parse_str::<syn::Path>(&format!("{}::{}", proto_path, method.output_type))
+            .unwrap()
+            .to_token_stream()
+    };
+
+    (request, response)
 }

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<T: #server_trait >(pub Arc<T>);
@@ -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<T: #server_trait >(pub Arc<T>);
@@ -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);
 

tonic-build/tests/protos/empty.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);
+}

tonic-build/tests/protos/wellknown.proto

@@ -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();
 }