use crate::generate_doc_comments; use proc_macro2::TokenStream; use prost_build::{Method, Service}; use quote::{format_ident, quote}; use syn::Path; pub(crate) fn generate(service: &Service, proto: &str) -> TokenStream { let service_ident = quote::format_ident!("{}Client", service.name); let methods = generate_methods(service, proto); let connect = generate_connect(&service_ident); let service_doc = generate_doc_comments(&service.comments.leading); quote! { #service_doc pub struct #service_ident { inner: tonic::client::Grpc, } #connect impl #service_ident where T: tonic::client::GrpcService, T::ResponseBody: Body + HttpBody + Send + 'static, T::Error: Into, ::Error: Into + Send, ::Data: Into + Send, { pub fn new(inner: T) -> Self { let inner = tonic::client::Grpc::new(inner); Self { inner } } pub async fn ready(&mut self) -> Result<(), tonic::Status> { self.inner.ready().await.map_err(|e| { tonic::Status::new(tonic::Code::Unknown, format!("Service was not ready: {}", e.into())) }) } #methods } impl Clone for #service_ident { fn clone(&self) -> Self { Self { inner: self.inner.clone(), } } } } } #[cfg(feature = "transport")] fn generate_connect(service_ident: &syn::Ident) -> TokenStream { quote! { impl #service_ident { pub fn connect(dst: D) -> Result where D: std::convert::TryInto, D::Error: Into, { tonic::transport::Channel::builder().build(dst).map(|c| Self::new(c)) } } } } #[cfg(not(feature = "transport"))] fn generate_connect() -> TokenStream { TokenStream::new() } fn generate_methods(service: &Service, proto: &str) -> TokenStream { let mut stream = TokenStream::new(); for method in &service.methods { let path = format!( "/{}.{}/{}", service.package, service.proto_name, method.proto_name ); stream.extend(generate_doc_comments(&method.comments.leading)); let method = match (method.client_streaming, method.server_streaming) { (false, false) => generate_unary(method, &proto, path), (false, true) => generate_server_streaming(method, &proto, path), (true, false) => generate_client_streaming(method, &proto, path), (true, true) => generate_streaming(method, &proto, path), }; stream.extend(method); } stream } fn generate_unary(method: &Method, proto: &str, path: String) -> TokenStream { let ident = format_ident!("{}", method.name); let request: Path = syn::parse_str(&format!("{}::{}", proto, method.input_type)).unwrap(); let response: Path = syn::parse_str(&format!("{}::{}", proto, method.output_type)).unwrap(); quote! { pub async fn #ident(&mut self, request: tonic::Request<#request>) -> Result, tonic::Status> { self.ready().await?; let codec = tonic::codec::ProstCodec::new(); let path = http::uri::PathAndQuery::from_static(#path); self.inner.unary(request, path, codec).await } } } fn generate_server_streaming(method: &Method, proto: &str, path: String) -> TokenStream { let ident = format_ident!("{}", method.name); let request: Path = syn::parse_str(&format!("{}::{}", proto, method.input_type)).unwrap(); let response: Path = syn::parse_str(&format!("{}::{}", proto, method.output_type)).unwrap(); quote! { pub async fn #ident(&mut self, request: tonic::Request<#request>) -> Result>, tonic::Status> { self.ready().await?; let codec = tonic::codec::ProstCodec::new(); let path = http::uri::PathAndQuery::from_static(#path); self.inner.server_streaming(request, path, codec).await } } } fn generate_client_streaming(method: &Method, proto: &str, path: String) -> TokenStream { let ident = format_ident!("{}", method.name); let request: Path = syn::parse_str(&format!("{}::{}", proto, method.input_type)).unwrap(); let response: Path = syn::parse_str(&format!("{}::{}", proto, method.output_type)).unwrap(); quote! { pub async fn #ident(&mut self, request: tonic::Request) -> Result, tonic::Status> where S: Stream> + Send + 'static, { self.ready().await?; let codec = tonic::codec::ProstCodec::new(); let path = http::uri::PathAndQuery::from_static(#path); self.inner.client_streaming(request, path, codec).await } } } fn generate_streaming(method: &Method, proto: &str, path: String) -> TokenStream { let ident = format_ident!("{}", method.name); let request: Path = syn::parse_str(&format!("{}::{}", proto, method.input_type)).unwrap(); let response: Path = syn::parse_str(&format!("{}::{}", proto, method.output_type)).unwrap(); quote! { pub async fn #ident(&mut self, request: tonic::Request) -> Result>, tonic::Status> where S: Stream> + Send + 'static, { self.ready().await?; let codec = tonic::codec::ProstCodec::new(); let path = http::uri::PathAndQuery::from_static(#path); self.inner.streaming(request, path, codec).await } } }