1 use std::collections::HashSet; 2 3 use proc_macro2::TokenStream; 4 5 use crate::{Attributes, Service}; 6 7 /// Builder for the generic code generation of server and clients. 8 #[derive(Debug)] 9 pub struct CodeGenBuilder { 10 emit_package: bool, 11 compile_well_known_types: bool, 12 attributes: Attributes, 13 build_transport: bool, 14 disable_comments: HashSet<String>, 15 use_arc_self: bool, 16 generate_default_stubs: bool, 17 } 18 19 impl CodeGenBuilder { 20 /// Create a new code gen builder with default options. new() -> Self21 pub fn new() -> Self { 22 Default::default() 23 } 24 25 /// Enable code generation to emit the package name. emit_package(&mut self, enable: bool) -> &mut Self26 pub fn emit_package(&mut self, enable: bool) -> &mut Self { 27 self.emit_package = enable; 28 self 29 } 30 31 /// Attributes that will be added to `mod` and `struct` items. 32 /// 33 /// Reference [`Attributes`] for more information. attributes(&mut self, attributes: Attributes) -> &mut Self34 pub fn attributes(&mut self, attributes: Attributes) -> &mut Self { 35 self.attributes = attributes; 36 self 37 } 38 39 /// Enable transport code to be generated, this requires `tonic`'s `transport` 40 /// feature. 41 /// 42 /// This allows codegen level control of generating the transport code and 43 /// is a work around when other crates in a workspace enable this feature. build_transport(&mut self, build_transport: bool) -> &mut Self44 pub fn build_transport(&mut self, build_transport: bool) -> &mut Self { 45 self.build_transport = build_transport; 46 self 47 } 48 49 /// Enable compiling well known types, this will force codegen to not 50 /// use the well known types from `prost-types`. compile_well_known_types(&mut self, enable: bool) -> &mut Self51 pub fn compile_well_known_types(&mut self, enable: bool) -> &mut Self { 52 self.compile_well_known_types = enable; 53 self 54 } 55 56 /// Disable comments based on a proto path. disable_comments(&mut self, disable_comments: HashSet<String>) -> &mut Self57 pub fn disable_comments(&mut self, disable_comments: HashSet<String>) -> &mut Self { 58 self.disable_comments = disable_comments; 59 self 60 } 61 62 /// Emit `Arc<Self>` instead of `&self` in service trait. use_arc_self(&mut self, enable: bool) -> &mut Self63 pub fn use_arc_self(&mut self, enable: bool) -> &mut Self { 64 self.use_arc_self = enable; 65 self 66 } 67 68 /// Enable or disable returning automatic unimplemented gRPC error code for generated traits. generate_default_stubs(&mut self, generate_default_stubs: bool) -> &mut Self69 pub fn generate_default_stubs(&mut self, generate_default_stubs: bool) -> &mut Self { 70 self.generate_default_stubs = generate_default_stubs; 71 self 72 } 73 74 /// Generate client code based on `Service`. 75 /// 76 /// This takes some `Service` and will generate a `TokenStream` that contains 77 /// a public module with the generated client. generate_client(&self, service: &impl Service, proto_path: &str) -> TokenStream78 pub fn generate_client(&self, service: &impl Service, proto_path: &str) -> TokenStream { 79 crate::client::generate_internal( 80 service, 81 self.emit_package, 82 proto_path, 83 self.compile_well_known_types, 84 self.build_transport, 85 &self.attributes, 86 &self.disable_comments, 87 ) 88 } 89 90 /// Generate server code based on `Service`. 91 /// 92 /// This takes some `Service` and will generate a `TokenStream` that contains 93 /// a public module with the generated client. generate_server(&self, service: &impl Service, proto_path: &str) -> TokenStream94 pub fn generate_server(&self, service: &impl Service, proto_path: &str) -> TokenStream { 95 crate::server::generate_internal( 96 service, 97 self.emit_package, 98 proto_path, 99 self.compile_well_known_types, 100 &self.attributes, 101 &self.disable_comments, 102 self.use_arc_self, 103 self.generate_default_stubs, 104 ) 105 } 106 } 107 108 impl Default for CodeGenBuilder { default() -> Self109 fn default() -> Self { 110 Self { 111 emit_package: true, 112 compile_well_known_types: false, 113 attributes: Attributes::default(), 114 build_transport: true, 115 disable_comments: HashSet::default(), 116 use_arc_self: false, 117 generate_default_stubs: false, 118 } 119 } 120 } 121