tonic_build/
code_gen.rs

1use std::collections::HashSet;
2
3use proc_macro2::TokenStream;
4
5use crate::{Attributes, Service};
6
7/// Builder for the generic code generation of server and clients.
8#[derive(Debug)]
9pub 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
19impl CodeGenBuilder {
20    /// Create a new code gen builder with default options.
21    pub fn new() -> Self {
22        Default::default()
23    }
24
25    /// Enable code generation to emit the package name.
26    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.
34    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.
44    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`.
51    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.
57    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.
63    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.
69    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.
78    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.
94    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
108impl Default for CodeGenBuilder {
109    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}