tower_grpc_build/
lib.rs

1#![doc(html_root_url = "https://docs.rs/tower-grpc-build/0.1.0")]
2#![deny(rust_2018_idioms)]
3#![cfg_attr(test, deny(warnings))]
4
5mod client;
6mod server;
7
8use heck::CamelCase;
9use std::io;
10use std::path::Path;
11
12/// Code generation configuration
13pub struct Config {
14    prost: prost_build::Config,
15    build_client: bool,
16    build_server: bool,
17}
18
19struct ServiceGenerator {
20    client: Option<client::ServiceGenerator>,
21    server: Option<server::ServiceGenerator>,
22    root_scope: codegen::Scope,
23}
24
25impl Config {
26    /// Returns a new `Config` with pre-configured prost.
27    ///
28    /// You can tweak the configuration how the proto buffers are generated and use this config.
29    pub fn from_prost(prost: prost_build::Config) -> Self {
30        Config {
31            prost,
32            // Enable client code gen by default
33            build_client: true,
34
35            // Disable server code gen by default
36            build_server: false,
37        }
38    }
39
40    /// Returns a new `Config` with default values.
41    pub fn new() -> Self {
42        Self::from_prost(prost_build::Config::new())
43    }
44
45    /// Enable gRPC client code generation
46    pub fn enable_client(&mut self, enable: bool) -> &mut Self {
47        self.build_client = enable;
48        self
49    }
50
51    /// Enable gRPC server code generation
52    pub fn enable_server(&mut self, enable: bool) -> &mut Self {
53        self.build_server = enable;
54        self
55    }
56
57    /// Generate code
58    pub fn build<P>(&mut self, protos: &[P], includes: &[P]) -> io::Result<()>
59    where
60        P: AsRef<Path>,
61    {
62        let client = if self.build_client {
63            Some(client::ServiceGenerator)
64        } else {
65            None
66        };
67        let server = if self.build_server {
68            Some(server::ServiceGenerator)
69        } else {
70            None
71        };
72
73        // Set or reset the service generator.
74        self.prost.service_generator(Box::new(ServiceGenerator {
75            client,
76            server,
77            root_scope: codegen::Scope::new(),
78        }));
79
80        self.prost.compile_protos(protos, includes)
81    }
82}
83
84impl prost_build::ServiceGenerator for ServiceGenerator {
85    fn generate(&mut self, service: prost_build::Service, _buf: &mut String) {
86        // Note that neither this implementation of `generate` nor the
87        // implementations for `client::ServiceGenerator` and
88        // `server::ServiceGenerator` will actually output any code to the
89        // buffer; all code is written out in the implementation of the
90        // `ServiceGenerator::finalize` function on this type.
91        if let Some(ref mut client_generator) = self.client {
92            client_generator.generate(&service, &mut self.root_scope);
93        }
94        if let Some(ref mut server_generator) = self.server {
95            server_generator.generate(&service, &mut self.root_scope);
96        }
97    }
98
99    fn finalize(&mut self, buf: &mut String) {
100        // Rather than outputting each service to the buffer as it's generated,
101        // we generate the code in our root `codegen::Scope`, which is shared
102        // between the generation of each service in the proto file. Unlike a
103        // string, codegen provides us with something not unlike a simplified
104        // Rust AST, making it easier for us to add new items to modules
105        // defined by previous service generator invocations. As we want to
106        // output the client and server implementations for each service in the
107        // proto file in one `client` or `server` module in the generated code,
108        // we wait until all the services have been generated before actually
109        // outputting to the buffer.
110        let mut fmt = codegen::Formatter::new(buf);
111        self.root_scope
112            .fmt(&mut fmt)
113            .expect("formatting root scope failed!");
114        // Reset the root scope so that the service generator is ready to
115        // generate another file. this prevents the code generated for *this*
116        // file being present in the next file.
117        self.root_scope = codegen::Scope::new();
118    }
119}
120
121/// Extension trait for importing generated types into `codegen::Scope`s.
122trait ImportType {
123    fn import_type(&mut self, ty: &str, level: usize);
124}
125
126// ===== impl ImportType =====
127
128impl ImportType for codegen::Scope {
129    fn import_type(&mut self, ty: &str, level: usize) {
130        if !is_imported_type(ty) && !is_native_type(ty) {
131            let (path, ty) = super_import(ty, level);
132
133            self.import(&path, &ty);
134        }
135    }
136}
137
138impl ImportType for codegen::Module {
139    fn import_type(&mut self, ty: &str, level: usize) {
140        self.scope().import_type(ty, level);
141    }
142}
143
144// ===== utility fns =====
145
146fn method_path(service: &prost_build::Service, method: &prost_build::Method) -> String {
147    format!(
148        "\"/{}.{}/{}\"",
149        service.package, service.proto_name, method.proto_name
150    )
151}
152
153fn lower_name(name: &str) -> String {
154    let mut ret = String::new();
155
156    for (i, ch) in name.chars().enumerate() {
157        if ch.is_uppercase() {
158            if i != 0 {
159                ret.push('_');
160            }
161
162            ret.push(ch.to_ascii_lowercase());
163        } else {
164            ret.push(ch);
165        }
166    }
167
168    ret
169}
170
171fn should_import(ty: &str) -> bool {
172    !is_imported_type(ty) && !is_native_type(ty)
173}
174
175fn is_imported_type(ty: &str) -> bool {
176    ty.split("::").map(|t| t == "super").next().unwrap()
177}
178
179fn is_native_type(ty: &str) -> bool {
180    match ty {
181        "()" => true,
182        _ => false,
183    }
184}
185
186fn super_import(ty: &str, level: usize) -> (String, String) {
187    let mut v: Vec<&str> = ty.split("::").collect();
188
189    assert!(!is_imported_type(ty));
190    assert!(!is_native_type(ty));
191
192    for _ in 0..level {
193        v.insert(0, "super");
194    }
195
196    let ty = v.pop().unwrap_or(ty);
197
198    (v.join("::"), ty.to_string())
199}
200
201fn unqualified(ty: &str, proto_ty: &str, level: usize) -> String {
202    if proto_ty == ".google.protobuf.Empty" {
203        return "()".to_string();
204    }
205
206    if !is_imported_type(ty) {
207        return ty.to_string();
208    }
209
210    let mut v: Vec<&str> = ty.split("::").collect();
211
212    for _ in 0..level {
213        v.insert(0, "super");
214    }
215
216    v.join("::")
217}
218
219/// Converts a `snake_case` identifier to an `UpperCamel` case Rust type
220/// identifier.
221///
222/// This is identical to the same [function] in `prost-build`, however, we
223/// reproduce it here as `prost` does not publically export it.
224///
225/// We need this as `prost-build` will only give us the snake-case transformed
226/// names for gRPC methods, but we need to use method names in types as well.
227///
228/// [function]: https://github.com/danburkert/prost/blob/d3b971ccd90df35d16069753d52289c0c85014e4/prost-build/src/ident.rs#L28-L38
229fn to_upper_camel(s: &str) -> String {
230    let mut ident = s.to_camel_case();
231
232    // Add a trailing underscore if the identifier matches a Rust keyword
233    // (https://doc.rust-lang.org/grammar.html#keywords).
234    if ident == "Self" {
235        ident.push('_');
236    }
237    ident
238}
239
240/// This function takes a `&prost_build::Comments` and formats them as a
241/// `String` to be used as a RustDoc comment.
242fn comments_to_rustdoc(comments: &prost_build::Comments) -> String {
243    comments
244        .leading
245        .iter()
246        .fold(String::new(), |acc, s| acc + s.trim_start() + "\n")
247}
tower_grpc_build/lib.rs

tower_grpc_build/
lib.rs
