rmqtt_protobuf_build/
lib.rs

1// Copyright 2019 PingCAP, Inc.
2
3//! Utility functions for generating Rust code from protobuf specifications.
4//!
5//! These functions panic liberally, they are designed to be used from build
6//! scripts, not in production.
7
8#[cfg(feature = "prost-codec")]
9mod wrapper;
10
11#[cfg(feature = "protobuf-codec")]
12mod protobuf_impl;
13
14#[cfg(feature = "prost-codec")]
15mod prost_impl;
16
17use bitflags::bitflags;
18use regex::Regex;
19use std::env;
20use std::env::var;
21use std::fmt::Write as _;
22use std::fs::{self, File};
23use std::io::Write;
24use std::path::{Path, PathBuf};
25use std::process::Command;
26use std::str::from_utf8;
27
28impl Builder {
29    /// Generate protobuf files using the appropriate codec.
30    ///
31    /// When both `prost-codec` and `protobuf-codec` are enabled simultaneously,
32    /// only the `prost-codec` generator runs (prost takes priority).
33    pub fn generate_files(&self) {
34        #[cfg(feature = "prost-codec")]
35        self.generate_files_prost();
36        #[cfg(all(feature = "protobuf-codec", not(feature = "prost-codec")))]
37        self.generate_files_protobuf();
38    }
39}
40
41// We use system protoc when its version matches,
42// otherwise use the protoc from bin which we bundle with the crate.
43fn get_protoc() -> String {
44    // $PROTOC overrides everything; if it isn't a useful version then fail.
45    if let Ok(s) = var("PROTOC") {
46        check_protoc_version(&s).expect("PROTOC version not usable");
47        return s;
48    }
49
50    if let Ok(s) = check_protoc_version("protoc") {
51        return s;
52    }
53
54    // The bundled protoc should always match the version
55    #[cfg(windows)]
56    {
57        let bin_path = Path::new(env!("CARGO_MANIFEST_DIR"))
58            .join("bin")
59            .join("protoc-win32.exe");
60        bin_path.display().to_string()
61    }
62
63    #[cfg(not(windows))]
64    protobuf_src::protoc().display().to_string()
65}
66
67fn check_protoc_version(protoc: &str) -> Result<String, ()> {
68    let ver_re = Regex::new(r"([0-9]+)\.([0-9]+)(\.[0-9])?").unwrap();
69    let output = Command::new(protoc).arg("--version").output();
70    match output {
71        Ok(o) => {
72            let caps = ver_re.captures(from_utf8(&o.stdout).unwrap()).unwrap();
73            let major = caps.get(1).unwrap().as_str().parse::<i16>().unwrap();
74            let minor = caps.get(2).unwrap().as_str().parse::<i16>().unwrap();
75            if (major, minor) >= (3, 1) {
76                return Ok(protoc.to_owned());
77            }
78            println!("The system `protoc` version mismatch, require >= 3.1.0, got {}.{}.x, fallback to the bundled `protoc`", major, minor);
79        }
80        Err(_) => println!("`protoc` not in PATH, try using the bundled protoc"),
81    };
82
83    Err(())
84}
85
86pub struct Builder {
87    files: Vec<String>,
88    includes: Vec<String>,
89    black_list: Vec<String>,
90    out_dir: String,
91    #[cfg(feature = "prost-codec")]
92    wrapper_opts: GenOpt,
93    package_name: Option<String>,
94    #[cfg(feature = "grpcio-protobuf-codec")]
95    re_export_services: bool,
96}
97
98impl Builder {
99    pub fn new() -> Builder {
100        Builder {
101            files: Vec::new(),
102            includes: vec!["include".to_owned(), "proto".to_owned()],
103            black_list: vec![
104                "protobuf".to_owned(),
105                "google".to_owned(),
106                "gogoproto".to_owned(),
107            ],
108            out_dir: format!("{}/protos", var("OUT_DIR").expect("No OUT_DIR defined")),
109            #[cfg(feature = "prost-codec")]
110            wrapper_opts: GenOpt::all(),
111            package_name: None,
112            #[cfg(feature = "grpcio-protobuf-codec")]
113            re_export_services: true,
114        }
115    }
116
117    pub fn include_google_protos(&mut self) -> &mut Self {
118        let path = format!("{}/include", env!("CARGO_MANIFEST_DIR"));
119        self.includes.push(path);
120        self
121    }
122
123    pub fn generate(&self) {
124        assert!(!self.files.is_empty(), "No files specified for generation");
125        self.prep_out_dir();
126        self.generate_files();
127        self.generate_mod_file();
128    }
129
130    /// This option is only used when generating Prost code. Otherwise, it is
131    /// silently ignored.
132    #[cfg(feature = "prost-codec")]
133    pub fn wrapper_options(&mut self, wrapper_opts: GenOpt) -> &mut Self {
134        self.wrapper_opts = wrapper_opts;
135        self
136    }
137
138    /// Finds proto files to operate on in the `proto_dir` directory.
139    pub fn search_dir_for_protos(&mut self, proto_dir: &str) -> &mut Self {
140        self.files = fs::read_dir(proto_dir)
141            .expect("Couldn't read proto directory")
142            .filter_map(|e| {
143                let e = e.expect("Couldn't list file");
144                if e.file_type().expect("File broken").is_dir() {
145                    None
146                } else {
147                    Some(format!("{}/{}", proto_dir, e.file_name().to_string_lossy()))
148                }
149            })
150            .collect();
151        self
152    }
153
154    pub fn files<T: ToString>(&mut self, files: &[T]) -> &mut Self {
155        self.files = files.iter().map(|t| t.to_string()).collect();
156        self
157    }
158
159    pub fn includes<T: ToString>(&mut self, includes: &[T]) -> &mut Self {
160        self.includes = includes.iter().map(|t| t.to_string()).collect();
161        self
162    }
163
164    pub fn append_include(&mut self, include: impl Into<String>) -> &mut Self {
165        self.includes.push(include.into());
166        self
167    }
168
169    pub fn black_list<T: ToString>(&mut self, black_list: &[T]) -> &mut Self {
170        self.black_list = black_list.iter().map(|t| t.to_string()).collect();
171        self
172    }
173
174    /// Add the name of an include file to the builder's black list.
175    ///
176    /// Files named on the black list are not made modules of the generated
177    /// program.
178    pub fn append_to_black_list(&mut self, include: impl Into<String>) -> &mut Self {
179        self.black_list.push(include.into());
180        self
181    }
182
183    pub fn out_dir(&mut self, out_dir: impl Into<String>) -> &mut Self {
184        self.out_dir = out_dir.into();
185        self
186    }
187
188    /// If specified, a module with the given name will be generated which re-exports
189    /// all generated items.
190    ///
191    /// This is ignored by Prost, since Prost uses the package names of protocols
192    /// in any case.
193    pub fn package_name(&mut self, package_name: impl Into<String>) -> &mut Self {
194        self.package_name = Some(package_name.into());
195        self
196    }
197
198    /// Whether services defined in separate modules should be re-exported from
199    /// their corresponding module. Default is `true`.
200    #[cfg(feature = "grpcio-protobuf-codec")]
201    pub fn re_export_services(&mut self, re_export_services: bool) -> &mut Self {
202        self.re_export_services = re_export_services;
203        self
204    }
205
206    fn generate_mod_file(&self) {
207        let mut f = File::create(format!("{}/mod.rs", self.out_dir)).unwrap();
208
209        let modules = self.list_rs_files().filter_map(|path| {
210            let name = path.file_stem().unwrap().to_str().unwrap();
211            if name.starts_with("wrapper_")
212                || name == "mod"
213                || self.black_list.iter().any(|i| name.contains(i))
214            {
215                return None;
216            }
217            Some((name.replace('-', "_"), name.to_owned()))
218        });
219
220        let mut exports = String::new();
221        for (module, file_name) in modules {
222            let wrapper_path = format!("{}/wrapper_{}.rs", self.out_dir, file_name);
223            let wrapper_exists = Path::new(&wrapper_path).exists();
224
225            if cfg!(not(feature = "prost-codec")) {
226                if wrapper_exists {
227                    // Wrapper exists (e.g., when both prost-codec and protobuf-codec
228                    // are enabled). Use nested module to include both files.
229                    writeln!(f, "pub mod {} {{", module).unwrap();
230                    writeln!(f, "    include!(\"{}.rs\");", file_name).unwrap();
231                    writeln!(f, "    include!(\"wrapper_{}.rs\");", file_name).unwrap();
232                    writeln!(f, "}}").unwrap();
233                    if self.package_name.is_some() {
234                        writeln!(exports, "pub use super::{}::*;", module).unwrap();
235                    }
236                } else if self.package_name.is_some() {
237                    writeln!(exports, "pub use super::{}::*;", module).unwrap();
238                    writeln!(f, "mod {};", module).unwrap();
239                } else {
240                    writeln!(f, "pub ").unwrap();
241                    writeln!(f, "mod {};", module).unwrap();
242                }
243                continue;
244            }
245
246            let mut level = 0;
247            for part in module.split('.') {
248                writeln!(f, "pub mod {} {{", part).unwrap();
249                level += 1;
250            }
251            writeln!(f, "include!(\"{}.rs\");", file_name,).unwrap();
252            if wrapper_exists {
253                writeln!(f, "include!(\"wrapper_{}.rs\");", file_name,).unwrap();
254            }
255            writeln!(f, "{}", "}\n".repeat(level)).unwrap();
256        }
257
258        if !exports.is_empty() {
259            writeln!(
260                f,
261                "pub mod {} {{ {} }}",
262                self.package_name.as_ref().unwrap(),
263                exports
264            )
265            .unwrap();
266        }
267    }
268
269    fn prep_out_dir(&self) {
270        if Path::new(&self.out_dir).exists() {
271            fs::remove_dir_all(&self.out_dir).unwrap();
272        }
273        fs::create_dir_all(&self.out_dir).unwrap();
274    }
275
276    // List all `.rs` files in `self.out_dir`.
277    fn list_rs_files(&self) -> impl Iterator<Item = PathBuf> {
278        fs::read_dir(&self.out_dir)
279            .expect("Couldn't read directory")
280            .filter_map(|e| {
281                let path = e.expect("Couldn't list file").path();
282                if path.extension() == Some(std::ffi::OsStr::new("rs")) {
283                    Some(path)
284                } else {
285                    None
286                }
287            })
288    }
289}
290
291impl Default for Builder {
292    fn default() -> Builder {
293        Builder::new()
294    }
295}
296
297bitflags! {
298    pub struct GenOpt: u32 {
299        /// Generate implementation for trait `::protobuf::Message`.
300        const MESSAGE = 0b0000_0001;
301        /// Generate getters.
302        const TRIVIAL_GET = 0b0000_0010;
303        /// Generate setters.
304        const TRIVIAL_SET = 0b0000_0100;
305        /// Generate the `new_` constructors.
306        const NEW = 0b0000_1000;
307        /// Generate `clear_*` functions.
308        const CLEAR = 0b0001_0000;
309        /// Generate `has_*` functions.
310        const HAS = 0b0010_0000;
311        /// Generate mutable getters.
312        const MUT = 0b0100_0000;
313        /// Generate `take_*` functions.
314        const TAKE = 0b1000_0000;
315        /// Except `impl protobuf::Message`.
316        const NO_MSG = Self::TRIVIAL_GET.bits
317         | Self::TRIVIAL_SET.bits
318         | Self::CLEAR.bits
319         | Self::HAS.bits
320         | Self::MUT.bits
321         | Self::TAKE.bits;
322        /// Except `new_` and `impl protobuf::Message`.
323        const ACCESSOR = Self::TRIVIAL_GET.bits
324         | Self::TRIVIAL_SET.bits
325         | Self::MUT.bits
326         | Self::TAKE.bits;
327    }
328}
rmqtt_protobuf_build/lib.rs

rmqtt_protobuf_build/
lib.rs
