Skip to content

Commit

Permalink
fix(codegen): Use wellknown types from prost-types (#49)
Browse files Browse the repository at this point in the history
This change introduces the ability to let prost-build
provide the fully qualifed type path for types that
start with `google.protobuf`.
  • Loading branch information
LucioFranco committed Oct 5, 2019
1 parent ed3e7e9 commit 4e1fcec
Show file tree
Hide file tree
Showing 6 changed files with 42 additions and 36 deletions.
12 changes: 4 additions & 8 deletions tonic-build/src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>)
Expand All @@ -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>)
Expand All @@ -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>)
Expand All @@ -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>)
Expand Down
27 changes: 18 additions & 9 deletions tonic-build/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
Expand Down Expand Up @@ -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)
}
15 changes: 5 additions & 10 deletions tonic-build/src/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down Expand Up @@ -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>);
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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>);
Expand Down Expand Up @@ -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);

Expand Down
7 changes: 0 additions & 7 deletions tonic-build/tests/protos/empty.proto

This file was deleted.

13 changes: 13 additions & 0 deletions tonic-build/tests/protos/wellknown.proto
Original file line number Diff line number Diff line change
@@ -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);
}
4 changes: 2 additions & 2 deletions tonic-build/tests/empty.rs → tonic-build/tests/wellknown.rs
Original file line number Diff line number Diff line change
@@ -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();
}

0 comments on commit 4e1fcec

Please sign in to comment.