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 } 16 17 impl CodeGenBuilder { 18 /// Create a new code gen builder with default options. 19 pub fn new() -> Self { 20 Default::default() 21 } 22 23 /// Enable code generation to emit the package name. 24 pub fn emit_package(&mut self, enable: bool) -> &mut Self { 25 self.emit_package = enable; 26 self 27 } 28 29 /// Attributes that will be added to `mod` and `struct` items. 30 /// 31 /// Reference [`Attributes`] for more information. 32 pub fn attributes(&mut self, attributes: Attributes) -> &mut Self { 33 self.attributes = attributes; 34 self 35 } 36 37 /// Enable transport code to be generated, this requires `tonic`'s `transport` 38 /// feature. 39 /// 40 /// This allows codegen level control of generating the transport code and 41 /// is a work around when other crates in a workspace enable this feature. 42 pub fn build_transport(&mut self, build_transport: bool) -> &mut Self { 43 self.build_transport = build_transport; 44 self 45 } 46 47 /// Enable compiling well knonw types, this will force codegen to not 48 /// use the well known types from `prost-types`. 49 pub fn compile_well_known_types(&mut self, enable: bool) -> &mut Self { 50 self.compile_well_known_types = enable; 51 self 52 } 53 54 /// Disable comments based on a proto path. 55 pub fn disable_comments(&mut self, disable_comments: HashSet<String>) -> &mut Self { 56 self.disable_comments = disable_comments; 57 self 58 } 59 60 /// Generate client code based on `Service`. 61 /// 62 /// This takes some `Service` and will generate a `TokenStream` that contains 63 /// a public module with the generated client. 64 pub fn generate_client(&self, service: &impl Service, proto_path: &str) -> TokenStream { 65 crate::client::generate_internal( 66 service, 67 self.emit_package, 68 proto_path, 69 self.compile_well_known_types, 70 self.build_transport, 71 &self.attributes, 72 &self.disable_comments, 73 ) 74 } 75 76 /// Generate server code based on `Service`. 77 /// 78 /// This takes some `Service` and will generate a `TokenStream` that contains 79 /// a public module with the generated client. 80 pub fn generate_server(&self, service: &impl Service, proto_path: &str) -> TokenStream { 81 crate::server::generate_internal( 82 service, 83 self.emit_package, 84 proto_path, 85 self.compile_well_known_types, 86 &self.attributes, 87 &self.disable_comments, 88 ) 89 } 90 } 91 92 impl Default for CodeGenBuilder { 93 fn default() -> Self { 94 Self { 95 emit_package: true, 96 compile_well_known_types: false, 97 attributes: Attributes::default(), 98 build_transport: true, 99 disable_comments: HashSet::default(), 100 } 101 } 102 } 103