Skip to content

Commit 363ae8a

Browse files
committed
build: refactor server trait gen to dedup some code
1 parent 14c1d91 commit 363ae8a

File tree

1 file changed

+44
-82
lines changed

1 file changed

+44
-82
lines changed

tonic-build/src/server.rs

Lines changed: 44 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -257,92 +257,54 @@ fn generate_trait_methods<T: Service>(
257257
quote!(&self)
258258
};
259259

260-
let method = match (
261-
method.client_streaming(),
262-
method.server_streaming(),
263-
generate_default_stubs,
264-
) {
265-
(false, false, true) => {
266-
quote! {
267-
#method_doc
268-
async fn #name(#self_param, request: tonic::Request<#req_message>)
269-
-> std::result::Result<tonic::Response<#res_message>, tonic::Status> {
270-
Err(tonic::Status::unimplemented("Not yet implemented"))
271-
}
272-
}
273-
}
274-
(false, false, false) => {
275-
quote! {
276-
#method_doc
277-
async fn #name(#self_param, request: tonic::Request<#req_message>)
278-
-> std::result::Result<tonic::Response<#res_message>, tonic::Status>;
279-
}
280-
}
281-
(true, false, true) => {
282-
quote! {
283-
#method_doc
284-
async fn #name(#self_param, request: tonic::Request<tonic::Streaming<#req_message>>)
285-
-> std::result::Result<tonic::Response<#res_message>, tonic::Status> {
286-
Err(tonic::Status::unimplemented("Not yet implemented"))
287-
}
288-
}
289-
}
290-
(true, false, false) => {
291-
quote! {
292-
#method_doc
293-
async fn #name(#self_param, request: tonic::Request<tonic::Streaming<#req_message>>)
294-
-> std::result::Result<tonic::Response<#res_message>, tonic::Status>;
295-
}
296-
}
297-
(false, true, true) => {
298-
quote! {
299-
#method_doc
300-
async fn #name(#self_param, request: tonic::Request<#req_message>)
301-
-> std::result::Result<tonic::Response<BoxStream<#res_message>>, tonic::Status> {
302-
Err(tonic::Status::unimplemented("Not yet implemented"))
303-
}
260+
let req_param_type = if method.client_streaming() {
261+
quote!(tonic::Request<tonic::Streaming<#req_message>>)
262+
} else {
263+
quote!(tonic::Request<#req_message>)
264+
};
265+
266+
let partial_sig = quote! {
267+
#method_doc
268+
async fn #name(#self_param, request: #req_param_type)
269+
};
270+
271+
let body_or_semicolon = if generate_default_stubs {
272+
quote! {
273+
{
274+
Err(tonic::Status::unimplemented("Not yet implemented"))
304275
}
305276
}
306-
(false, true, false) => {
307-
let stream = quote::format_ident!("{}Stream", method.identifier());
308-
let stream_doc = generate_doc_comment(format!(
309-
" Server streaming response type for the {} method.",
310-
method.identifier()
311-
));
312-
313-
quote! {
314-
#stream_doc
315-
type #stream: tonic::codegen::tokio_stream::Stream<Item = std::result::Result<#res_message, tonic::Status>> + std::marker::Send + 'static;
316-
317-
#method_doc
318-
async fn #name(#self_param, request: tonic::Request<#req_message>)
319-
-> std::result::Result<tonic::Response<Self::#stream>, tonic::Status>;
320-
}
277+
} else {
278+
quote!(;)
279+
};
280+
281+
let result = |ok| quote!(std::result::Result<#ok, tonic::Status>);
282+
let response_result = |message| result(quote!(tonic::Response<#message>));
283+
284+
let method = if !method.server_streaming() {
285+
let return_ty = response_result(res_message);
286+
quote! {
287+
#partial_sig -> #return_ty #body_or_semicolon
321288
}
322-
(true, true, true) => {
323-
quote! {
324-
#method_doc
325-
async fn #name(#self_param, request: tonic::Request<tonic::Streaming<#req_message>>)
326-
-> std::result::Result<tonic::Response<BoxStream<#res_message>>, tonic::Status> {
327-
Err(tonic::Status::unimplemented("Not yet implemented"))
328-
}
329-
}
289+
} else if generate_default_stubs {
290+
let return_ty = response_result(quote!(BoxStream<#res_message>));
291+
quote! {
292+
#partial_sig -> #return_ty #body_or_semicolon
330293
}
331-
(true, true, false) => {
332-
let stream = quote::format_ident!("{}Stream", method.identifier());
333-
let stream_doc = generate_doc_comment(format!(
334-
" Server streaming response type for the {} method.",
335-
method.identifier()
336-
));
337-
338-
quote! {
339-
#stream_doc
340-
type #stream: tonic::codegen::tokio_stream::Stream<Item = std::result::Result<#res_message, tonic::Status>> + std::marker::Send + 'static;
341-
342-
#method_doc
343-
async fn #name(#self_param, request: tonic::Request<tonic::Streaming<#req_message>>)
344-
-> std::result::Result<tonic::Response<Self::#stream>, tonic::Status>;
345-
}
294+
} else {
295+
let stream = quote::format_ident!("{}Stream", method.identifier());
296+
let stream_doc = generate_doc_comment(format!(
297+
" Server streaming response type for the {} method.",
298+
method.identifier()
299+
));
300+
let stream_item_ty = result(res_message);
301+
let stream_ty = quote!(tonic::codegen::tokio_stream::Stream<Item = #stream_item_ty> + std::marker::Send + 'static);
302+
let return_ty = response_result(quote!(Self::#stream));
303+
quote! {
304+
#stream_doc
305+
type #stream: #stream_ty;
306+
307+
#partial_sig -> #return_ty #body_or_semicolon
346308
}
347309
};
348310

0 commit comments

Comments
 (0)