tonic_build/
code_gen.rs

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