Nest server and client in their own mods

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<path::Path>,
 {
     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<P: AsRef<Path>>(
+    protos: &[P],
+    includes: &[P],
+    package: &str,
+    out_dir: impl AsRef<Path>,
+) -> 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);
+        }
     }
 }

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<T: #server_trait >(pub std::sync::Arc<T>);
+        struct #service_ident<T: #server_trait>(pub std::sync::Arc<T>);
 
-        impl<T: #server_trait > tonic::server::StreamingService<#request> for #service_ident <T>
+        impl<T: #server_trait> tonic::server::StreamingService<#request> for #service_ident <T>
         {
             type Response = #response;
             type ResponseStream = T::#response_stream;