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