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); quote! { pub struct #service_ident { inner: tonic::client::Grpc, } 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(), } } } } } 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 ); 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 } } }