chore(build): Refactor codegen traits (#302)

LucioFranco · web-flow · commit 2e082f8b3eea · 2020-03-29T12:57:11.000-04:00
Signed-off-by: Lucio Franco &lt;luciofranco14@gmail.com&gt;
diff --git a/examples/build.rs b/examples/build.rs
@@ -1,6 +1,6 @@
 fn main() {
-    tonic_build::prost::compile_protos("proto/helloworld/helloworld.proto").unwrap();
-    tonic_build::prost::compile_protos("proto/routeguide/route_guide.proto").unwrap();
-    tonic_build::prost::compile_protos("proto/echo/echo.proto").unwrap();
-    tonic_build::prost::compile_protos("proto/google/pubsub/pubsub.proto").unwrap();
+    tonic_build::compile_protos("proto/helloworld/helloworld.proto").unwrap();
+    tonic_build::compile_protos("proto/routeguide/route_guide.proto").unwrap();
+    tonic_build::compile_protos("proto/echo/echo.proto").unwrap();
+    tonic_build::compile_protos("proto/google/pubsub/pubsub.proto").unwrap();
 }
diff --git a/examples/routeguide-tutorial.md b/examples/routeguide-tutorial.md
@@ -823,7 +823,7 @@ opposed to at build time, placing the resulting modules wherever we need them.
 
 ```rust
 fn main() {
-    tonic_build::prost::configure()
+    tonic_build::configure()
         .build_client(false)
         .out_dir("another_crate/src/pb")
         .compile(&["path/my_proto.proto"], &["path"])
diff --git a/interop/build.rs b/interop/build.rs
@@ -1,7 +1,7 @@
 fn main() {
     let proto = "proto/grpc/testing/test.proto";
 
-    tonic_build::prost::compile_protos(proto).unwrap();
+    tonic_build::compile_protos(proto).unwrap();
 
     // prevent needing to rebuild if files (or deps) haven't changed
     println!("cargo:rerun-if-changed={}", proto);
diff --git a/tests/extern_path/my_application/build.rs b/tests/extern_path/my_application/build.rs
@@ -1,5 +1,5 @@
 fn main() -> Result<(), std::io::Error> {
-    tonic_build::prost::configure()
+    tonic_build::configure()
         .build_server(false)
         .build_client(true)
         .extern_path(".uuid", "::uuid")
diff --git a/tests/included_service/build.rs b/tests/included_service/build.rs
@@ -1,3 +1,3 @@
 fn main() {
-    tonic_build::prost::compile_protos("proto/includer.proto").unwrap();
+    tonic_build::compile_protos("proto/includer.proto").unwrap();
 }
diff --git a/tests/same_name/build.rs b/tests/same_name/build.rs
@@ -1,3 +1,3 @@
 fn main() {
-    tonic_build::prost::compile_protos("proto/foo.proto").unwrap();
+    tonic_build::compile_protos("proto/foo.proto").unwrap();
 }
diff --git a/tests/wellknown/build.rs b/tests/wellknown/build.rs
@@ -1,3 +1,3 @@
 fn main() {
-    tonic_build::prost::compile_protos("proto/wellknown.proto").unwrap();
+    tonic_build::compile_protos("proto/wellknown.proto").unwrap();
 }
diff --git a/tonic-build/README.md b/tonic-build/README.md
@@ -23,7 +23,7 @@ tonic-build = <tonic-version>
 
 ```rust
 fn main() -> Result<(), Box<dyn std::error::Error>> {
-    tonic_build::prost::compile_protos("proto/service.proto")?;
+    tonic_build::compile_protos("proto/service.proto")?;
     Ok(())
 }
 ```
@@ -32,7 +32,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
 ```rust
 fn main() -> Result<(), Box<dyn std::error::Error>> {
-   tonic_build::prost::configure()
+   tonic_build::configure()
         .build_server(false)
         .compile(
             &["proto/helloworld/helloworld.proto"],
diff --git a/tonic-build/src/client.rs b/tonic-build/src/client.rs
@@ -1,13 +1,13 @@
-use super::schema::{Context, Method, Service};
+use super::schema::{Method, Service};
 use crate::{generate_doc_comments, naive_snake_case};
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 
 /// Generate service for client
-pub fn generate<'a, T: Service<'a>>(service: &'a T, context: &T::Context) -> TokenStream {
+pub fn generate<T: Service>(service: &T, proto_path: &str) -> TokenStream {
     let service_ident = quote::format_ident!("{}Client", service.name());
     let client_mod = quote::format_ident!("{}_client", naive_snake_case(&service.name()));
-    let methods = generate_methods(service, context);
+    let methods = generate_methods(service, proto_path);
 
     let connect = generate_connect(&service_ident);
     let service_doc = generate_doc_comments(service.comment());
@@ -76,7 +76,7 @@ fn generate_connect(_service_ident: &syn::Ident) -> TokenStream {
     TokenStream::new()
 }
 
-fn generate_methods<'a, T: Service<'a>>(service: &'a T, context: &T::Context) -> TokenStream {
+fn generate_methods<T: Service>(service: &T, proto_path: &str) -> TokenStream {
     let mut stream = TokenStream::new();
 
     for method in service.methods() {
@@ -92,10 +92,10 @@ fn generate_methods<'a, T: Service<'a>>(service: &'a T, context: &T::Context) ->
         stream.extend(generate_doc_comments(method.comment()));
 
         let method = match (method.client_streaming(), method.server_streaming()) {
-            (false, false) => generate_unary(method, &context, path),
-            (false, true) => generate_server_streaming(method, &context, path),
-            (true, false) => generate_client_streaming(method, &context, path),
-            (true, true) => generate_streaming(method, &context, path),
+            (false, false) => generate_unary(method, proto_path, path),
+            (false, true) => generate_server_streaming(method, proto_path, path),
+            (true, false) => generate_client_streaming(method, proto_path, path),
+            (true, true) => generate_streaming(method, proto_path, path),
         };
 
         stream.extend(method);
@@ -104,14 +104,10 @@ fn generate_methods<'a, T: Service<'a>>(service: &'a T, context: &T::Context) ->
     stream
 }
 
-fn generate_unary<'a, T: Method<'a>>(
-    method: &T,
-    context: &T::Context,
-    path: String,
-) -> TokenStream {
-    let codec_name = syn::parse_str::<syn::Path>(context.codec_name()).unwrap();
+fn generate_unary<T: Method>(method: &T, proto_path: &str, path: String) -> TokenStream {
+    let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap();
     let ident = format_ident!("{}", method.name());
-    let (request, response) = method.request_response_name(context);
+    let (request, response) = method.request_response_name(proto_path);
 
     quote! {
         pub async fn #ident(
@@ -128,15 +124,11 @@ fn generate_unary<'a, T: Method<'a>>(
     }
 }
 
-fn generate_server_streaming<'a, T: Method<'a>>(
-    method: &T,
-    context: &T::Context,
-    path: String,
-) -> TokenStream {
-    let codec_name = syn::parse_str::<syn::Path>(context.codec_name()).unwrap();
+fn generate_server_streaming<T: Method>(method: &T, proto_path: &str, path: String) -> TokenStream {
+    let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap();
     let ident = format_ident!("{}", method.name());
 
-    let (request, response) = method.request_response_name(context);
+    let (request, response) = method.request_response_name(proto_path);
 
     quote! {
         pub async fn #ident(
@@ -153,15 +145,11 @@ fn generate_server_streaming<'a, T: Method<'a>>(
     }
 }
 
-fn generate_client_streaming<'a, T: Method<'a>>(
-    method: &T,
-    context: &T::Context,
-    path: String,
-) -> TokenStream {
-    let codec_name = syn::parse_str::<syn::Path>(context.codec_name()).unwrap();
+fn generate_client_streaming<T: Method>(method: &T, proto_path: &str, path: String) -> TokenStream {
+    let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap();
     let ident = format_ident!("{}", method.name());
 
-    let (request, response) = method.request_response_name(context);
+    let (request, response) = method.request_response_name(proto_path);
 
     quote! {
         pub async fn #ident(
@@ -178,15 +166,11 @@ fn generate_client_streaming<'a, T: Method<'a>>(
     }
 }
 
-fn generate_streaming<'a, T: Method<'a>>(
-    method: &T,
-    context: &T::Context,
-    path: String,
-) -> TokenStream {
-    let codec_name = syn::parse_str::<syn::Path>(context.codec_name()).unwrap();
+fn generate_streaming<T: Method>(method: &T, proto_path: &str, path: String) -> TokenStream {
+    let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap();
     let ident = format_ident!("{}", method.name());
 
-    let (request, response) = method.request_response_name(context);
+    let (request, response) = method.request_response_name(proto_path);
 
     quote! {
         pub async fn #ident(
diff --git a/tonic-build/src/lib.rs b/tonic-build/src/lib.rs
@@ -23,7 +23,7 @@
 //!
 //! ```rust,no_run
 //! fn main() -> Result<(), Box<dyn std::error::Error>> {
-//!     tonic_build::prost::compile_protos("proto/service.proto")?;
+//!     tonic_build::compile_protos("proto/service.proto")?;
 //!     Ok(())
 //! }
 //! ```
@@ -32,7 +32,7 @@
 //!
 //! ```rust,no_run
 //! fn main() -> Result<(), Box<dyn std::error::Error>> {
-//!    tonic_build::prost::configure()
+//!    tonic_build::configure()
 //!         .build_server(false)
 //!         .compile(
 //!             &["proto/helloworld/helloworld.proto"],
@@ -61,7 +61,11 @@ use quote::TokenStreamExt;
 
 /// Prost generator
 #[cfg(feature = "prost")]
-pub mod prost;
+mod prost;
+
+#[cfg(feature = "prost")]
+pub use prost::{compile_protos, configure, Builder};
+
 /// Traits to describe schema
 pub mod schema;
 
diff --git a/tonic-build/src/prost.rs b/tonic-build/src/prost.rs
@@ -5,32 +5,20 @@ use quote::ToTokens;
 use std::io;
 use std::path::{Path, PathBuf};
 
-impl<'a> schema::Commentable<'a> for Service {
+const PROST_CODEC_PATH: &'static str = "tonic::codec::ProstCodec";
+
+impl schema::Commentable for Service {
     type Comment = String;
-    type CommentContainer = &'a Vec<Self::Comment>;
 
-    fn comment(&'a self) -> Self::CommentContainer {
-        &self.comments.leading
+    fn comment(&self) -> &[Self::Comment] {
+        &self.comments.leading[..]
     }
 }
 
-/// Context data used while generate prost service
-#[derive(Debug)]
-pub struct ProstContext {
-    /// relative path to proto definitions from service definitions
-    pub proto_path: String,
-}
-
-impl schema::Context for ProstContext {
-    fn codec_name(&self) -> &str {
-        "tonic::codec::ProstCodec"
-    }
-}
+impl schema::Service for Service {
+    const CODEC_PATH: &'static str = PROST_CODEC_PATH;
 
-impl<'a> schema::Service<'a> for Service {
     type Method = Method;
-    type MethodContainer = &'a Vec<Self::Method>;
-    type Context = ProstContext;
 
     fn name(&self) -> &str {
         &self.name
@@ -44,22 +32,21 @@ impl<'a> schema::Service<'a> for Service {
         &self.proto_name
     }
 
-    fn methods(&'a self) -> Self::MethodContainer {
-        &self.methods
+    fn methods(&self) -> &[Self::Method] {
+        &self.methods[..]
     }
 }
 
-impl<'a> schema::Commentable<'a> for Method {
+impl schema::Commentable for Method {
     type Comment = String;
-    type CommentContainer = &'a Vec<Self::Comment>;
 
-    fn comment(&'a self) -> Self::CommentContainer {
-        &self.comments.leading
+    fn comment(&self) -> &[Self::Comment] {
+        &self.comments.leading[..]
     }
 }
 
-impl<'a> schema::Method<'a> for Method {
-    type Context = ProstContext;
+impl schema::Method for Method {
+    const CODEC_PATH: &'static str = PROST_CODEC_PATH;
 
     fn name(&self) -> &str {
         &self.name
@@ -77,13 +64,13 @@ impl<'a> schema::Method<'a> for Method {
         self.server_streaming
     }
 
-    fn request_response_name(&self, context: &Self::Context) -> (TokenStream, TokenStream) {
+    fn request_response_name(&self, proto_path: &str) -> (TokenStream, TokenStream) {
         let request = if self.input_proto_type.starts_with(".google.protobuf")
             || self.input_type.starts_with("::")
         {
             self.input_type.parse::<TokenStream>().unwrap()
         } else {
-            syn::parse_str::<syn::Path>(&format!("{}::{}", context.proto_path, self.input_type))
+            syn::parse_str::<syn::Path>(&format!("{}::{}", proto_path, self.input_type))
                 .unwrap()
                 .to_token_stream()
         };
@@ -93,7 +80,7 @@ impl<'a> schema::Method<'a> for Method {
         {
             self.output_type.parse::<TokenStream>().unwrap()
         } else {
-            syn::parse_str::<syn::Path>(&format!("{}::{}", context.proto_path, self.output_type))
+            syn::parse_str::<syn::Path>(&format!("{}::{}", proto_path, self.output_type))
                 .unwrap()
                 .to_token_stream()
         };
@@ -145,17 +132,13 @@ impl ServiceGenerator {
 
 impl prost_build::ServiceGenerator for ServiceGenerator {
     fn generate(&mut self, service: prost_build::Service, _buf: &mut String) {
-        let context = ProstContext {
-            proto_path: String::from("super"),
-        };
-
         if self.builder.build_server {
-            let server = server::generate(&service, &context);
+            let server = server::generate(&service, &self.builder.proto_path);
             self.servers.extend(server);
         }
 
         if self.builder.build_client {
-            let client = client::generate(&service, &context);
+            let client = client::generate(&service, &self.builder.proto_path);
             self.clients.extend(client);
         }
     }
@@ -197,6 +180,7 @@ pub struct Builder {
     pub(crate) extern_path: Vec<(String, String)>,
     pub(crate) field_attributes: Vec<(String, String)>,
     pub(crate) type_attributes: Vec<(String, String)>,
+    pub(crate) proto_path: String,
 
     out_dir: Option<PathBuf>,
     #[cfg(feature = "rustfmt")]
@@ -262,6 +246,15 @@ impl Builder {
         self
     }
 
+    /// Set the path to where tonic will search for the Request/Response proto structs
+    /// live relative to the module where you call `include_proto!`.
+    ///
+    /// This defaults to `super` since tonic will generate code in a module.
+    pub fn proto_path(mut self, proto_path: impl AsRef<str>) -> Self {
+        self.proto_path = proto_path.as_ref().to_string();
+        self
+    }
+
     /// Compile the .proto files and execute code generation.
     pub fn compile<P: AsRef<Path>>(self, protos: &[P], includes: &[P]) -> io::Result<()> {
         let out_dir = if let Some(out_dir) = self.out_dir.as_ref() {
@@ -297,6 +290,7 @@ pub fn configure() -> Builder {
         extern_path: Vec::new(),
         field_attributes: Vec::new(),
         type_attributes: Vec::new(),
+        proto_path: "super".to_string(),
         #[cfg(feature = "rustfmt")]
         format: true,
     }
@@ -306,8 +300,8 @@ pub fn configure() -> Builder {
 ///
 /// The include directory will be the parent folder of the specified path.
 /// The package name will be the filename without the extension.
-pub fn compile_protos(proto_path: impl AsRef<Path>) -> io::Result<()> {
-    let proto_path: &Path = proto_path.as_ref();
+pub fn compile_protos(proto: impl AsRef<Path>) -> io::Result<()> {
+    let proto_path: &Path = proto.as_ref();
 
     // directory the main .proto file resides in
     let proto_dir = proto_path
diff --git a/tonic-build/src/schema.rs b/tonic-build/src/schema.rs
diff --git a/tonic-build/src/server.rs b/tonic-build/src/server.rs

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,3 @@`
`1`	`1`	`fn main() {`
`2`		`- tonic_build::prost::compile_protos("proto/includer.proto").unwrap();`
	`2`	`+ tonic_build::compile_protos("proto/includer.proto").unwrap();`
`3`	`3`	`}`