diff --git a/tonic-build/src/lib.rs b/tonic-build/src/lib.rs index 3f6496e..a62c426 100644 --- a/tonic-build/src/lib.rs +++ b/tonic-build/src/lib.rs @@ -1,5 +1,6 @@ +use proc_macro2::TokenStream; use prost_build::Config; -use std::{io, path, process::Command}; +use std::{io, path, path::Path, process::Command}; mod client; mod service; @@ -9,13 +10,25 @@ where P: AsRef, { let out_dir = std::env::var("OUT_DIR").unwrap(); + compile_protos_with_out_dir(protos, includes, package, out_dir.as_str()) +} + +pub fn compile_protos_with_out_dir>( + protos: &[P], + includes: &[P], + package: &str, + out_dir: impl AsRef, +) -> io::Result<()> { let mut config = Config::new(); - config.service_generator(Box::new(ServiceGenerator {})); - config.out_dir(&out_dir); + config.service_generator(Box::new(ServiceGenerator::default())); + config.out_dir(out_dir.as_ref()); config.compile_protos(protos, includes)?; - fmt(&out_dir, &format!("{}.rs", package)); + fmt( + out_dir.as_ref().to_str().expect("Execpted utf8 out_dir"), + &format!("{}.rs", package), + ); Ok(()) } @@ -34,17 +47,44 @@ fn fmt(out_dir: &str, file: &str) { assert!(out.status.success()); } -pub struct ServiceGenerator {} +#[derive(Default)] +pub struct ServiceGenerator { + clients: TokenStream, + servers: TokenStream, +} impl prost_build::ServiceGenerator for ServiceGenerator { - fn generate(&mut self, service: prost_build::Service, buf: &mut String) { - let path = "self"; + fn generate(&mut self, service: prost_build::Service, _buf: &mut String) { + let path = "super"; + let server = service::generate(&service, path); - let code = format!("{}", server); - buf.push_str(&code); + self.servers.extend(server); let client = client::generate(&service, path); - let code = format!("{}", client); - buf.push_str(&code); + self.clients.extend(client); + } + + fn finalize(&mut self, buf: &mut String) { + if !self.clients.is_empty() && !self.servers.is_empty() { + let clients = &self.clients; + let servers = &self.servers; + + let service = quote::quote! { + pub mod client { + #![allow(unused_variables, dead_code, missing_docs)] + + #clients + } + + pub mod server { + #![allow(unused_variables, dead_code, missing_docs)] + + #servers + } + }; + + let code = format!("{}", service); + buf.push_str(&code); + } } } diff --git a/tonic-build/src/service.rs b/tonic-build/src/service.rs index 081bfa6..4829bdb 100644 --- a/tonic-build/src/service.rs +++ b/tonic-build/src/service.rs @@ -321,12 +321,10 @@ fn generate_streaming( let response_stream = quote::format_ident!("{}Stream", method.proto_name); - // TODO: parse response stream type, if it is a concrete type then use that - // as the ResponseStream type, if it is a impl Trait then we need to box. quote! { - struct #service_ident(pub std::sync::Arc); + struct #service_ident(pub std::sync::Arc); - impl tonic::server::StreamingService<#request> for #service_ident + impl tonic::server::StreamingService<#request> for #service_ident { type Response = #response; type ResponseStream = T::#response_stream; diff --git a/tonic-examples/src/helloworld/client.rs b/tonic-examples/src/helloworld/client.rs index 4ad3fae..fd673ec 100644 --- a/tonic-examples/src/helloworld/client.rs +++ b/tonic-examples/src/helloworld/client.rs @@ -4,15 +4,17 @@ pub mod hello_world { include!(concat!(env!("OUT_DIR"), "/helloworld.rs")); } +use hello_world::{client::GreeterClient, HelloRequest}; + #[tokio::main] async fn main() -> Result<(), Box> { let origin = vec![http::Uri::from_static("http://[::1]:50051").into()]; let svc = Channel::builder().balance_list(origin)?; - let mut client = hello_world::GreeterClient::new(svc); + let mut client = GreeterClient::new(svc); - let request = tonic::Request::new(hello_world::HelloRequest { + let request = tonic::Request::new(HelloRequest { name: "hello".into(), }); diff --git a/tonic-examples/src/helloworld/server.rs b/tonic-examples/src/helloworld/server.rs index c7e464b..bc7f256 100644 --- a/tonic-examples/src/helloworld/server.rs +++ b/tonic-examples/src/helloworld/server.rs @@ -4,17 +4,22 @@ pub mod hello_world { include!(concat!(env!("OUT_DIR"), "/helloworld.rs")); } +use hello_world::{ + server::{Greeter, GreeterServer}, + HelloReply, HelloRequest, +}; + #[derive(Default)] pub struct MyGreeter { data: String, } #[tonic::async_trait] -impl hello_world::Greeter for MyGreeter { +impl Greeter for MyGreeter { async fn say_hello( &self, - request: Request, - ) -> Result, Status> { + request: Request, + ) -> Result, Status> { println!("Got a request: {:?}", request); let string = &self.data; @@ -34,7 +39,7 @@ async fn main() -> Result<(), Box> { let greeter = MyGreeter::default(); Server::builder() - .serve(addr, hello_world::GreeterServer::new(greeter)) + .serve(addr, GreeterServer::new(greeter)) .await?; Ok(()) diff --git a/tonic-examples/src/routeguide/client.rs b/tonic-examples/src/routeguide/client.rs index 455a152..1fd3d3f 100644 --- a/tonic-examples/src/routeguide/client.rs +++ b/tonic-examples/src/routeguide/client.rs @@ -8,12 +8,14 @@ mod route_guide { include!(concat!(env!("OUT_DIR"), "/routeguide.rs")); } +use route_guide::client::RouteGuideClient; + #[tokio::main] async fn main() -> Result<(), Box> { let origin = http::Uri::from_static("http://[::1]:10000"); let svc = Channel::builder().build(origin)?; - let mut client = route_guide::RouteGuideClient::new(svc); + let mut client = RouteGuideClient::new(svc); let start = Instant::now(); diff --git a/tonic-examples/src/routeguide/server.rs b/tonic-examples/src/routeguide/server.rs index cb657e7..24a4c61 100644 --- a/tonic-examples/src/routeguide/server.rs +++ b/tonic-examples/src/routeguide/server.rs @@ -14,7 +14,7 @@ pub mod routeguide { include!(concat!(env!("OUT_DIR"), "/routeguide.rs")); } -use routeguide::{Feature, Point, Rectangle, RouteNote, RouteSummary}; +use routeguide::{server, Feature, Point, Rectangle, RouteNote, RouteSummary}; #[derive(Debug)] pub struct RouteGuide { @@ -28,7 +28,7 @@ struct State { } #[tonic::async_trait] -impl routeguide::RouteGuide for RouteGuide { +impl server::RouteGuide for RouteGuide { async fn get_feature(&self, request: Request) -> Result, Status> { println!("GetFeature = {:?}", request); @@ -168,7 +168,7 @@ async fn main() -> Result<(), Box> { }, }; - let svc = routeguide::RouteGuideServer::new(route_guide); + let svc = server::RouteGuideServer::new(route_guide); Server::builder().serve(addr, svc).await?; diff --git a/tonic-interop/src/client.rs b/tonic-interop/src/client.rs index b2e1204..c460478 100644 --- a/tonic-interop/src/client.rs +++ b/tonic-interop/src/client.rs @@ -1,4 +1,4 @@ -use crate::{pb::*, test_assert, TestAssertion}; +use crate::{pb::client::*, pb::*, test_assert, TestAssertion}; use futures_util::{future, stream, SinkExt, StreamExt}; use tokio::sync::mpsc; use tonic::transport::Channel; diff --git a/tonic-interop/src/server.rs b/tonic-interop/src/server.rs index 9bf81fa..ec914a4 100644 --- a/tonic-interop/src/server.rs +++ b/tonic-interop/src/server.rs @@ -1,9 +1,9 @@ -use crate::pb::*; +use crate::pb::{self, *}; use std::pin::Pin; use tonic::{Code, Request, Response, Status}; -pub fn create() -> TestServiceServer { - TestServiceServer::new(TestService { +pub fn create() -> pb::server::TestServiceServer { + server::TestServiceServer::new(TestService { data: String::new(), }) } @@ -19,7 +19,7 @@ type Stream = Pin> + Send + 'static>>; #[tonic::async_trait] -impl crate::pb::TestService for TestService { +impl pb::server::TestService for TestService { async fn empty_call(&self, _request: Request) -> Result { println!("empty_call"); Ok(Response::new(Empty {}))