protobuf_build/
protobuf_impl.rs

1// Copyright 2019 PingCAP, Inc.
2
3use std::fs::File;
4use std::io::{Read, Write};
5use std::path::Path;
6use std::process::Command;
7
8use protobuf::Message;
9use regex::Regex;
10
11use crate::get_protoc;
12use crate::Builder;
13
14impl Builder {
15    pub fn generate_files(&self) {
16        let mut cmd = Command::new(get_protoc());
17        let desc_file = format!("{}/mod.desc", self.out_dir);
18        for i in &self.includes {
19            cmd.arg(format!("-I{}", i));
20        }
21        cmd.arg("--include_imports")
22            .arg("--include_source_info")
23            .arg("-o")
24            .arg(&desc_file);
25        for f in &self.files {
26            cmd.arg(f);
27        }
28        println!("executing {:?}", cmd);
29        match cmd.status() {
30            Ok(e) if e.success() => {}
31            e => panic!("failed to generate descriptor set files: {:?}", e),
32        }
33
34        let desc_bytes = std::fs::read(&desc_file).unwrap();
35        let mut desc = protobuf::descriptor::FileDescriptorSet::new();
36        desc.merge_from_bytes(&desc_bytes).unwrap();
37        desc.check_initialized().unwrap();
38
39        let mut files_to_generate = Vec::new();
40        'outer: for file in &self.files {
41            for include in &self.includes {
42                if let Ok(truncated) = Path::new(file).strip_prefix(include) {
43                    files_to_generate.push(format!("{}", truncated.display()));
44                    continue 'outer;
45                }
46            }
47
48            panic!(
49                "file {:?} is not found in includes {:?}",
50                file, self.includes
51            );
52        }
53
54        protobuf_codegen::gen_and_write(
55            desc.get_file(),
56            &files_to_generate,
57            Path::new(&self.out_dir),
58            &protobuf_codegen::Customize::default(),
59        )
60        .unwrap();
61        self.generate_grpcio(desc.get_file(), &files_to_generate);
62        self.import_grpcio();
63        self.replace_read_unknown_fields();
64    }
65
66    /// Convert protobuf files to use the old way of reading protobuf enums.
67    // FIXME: Remove this once stepancheg/rust-protobuf#233 is resolved.
68    fn replace_read_unknown_fields(&self) {
69        let regex =
70            Regex::new(r"::protobuf::rt::read_proto3_enum_with_unknown_fields_into\(([^,]+), ([^,]+), &mut ([^,]+), [^\)]+\)\?").unwrap();
71        self.list_rs_files().for_each(|path| {
72            let mut text = String::new();
73            let mut f = File::open(&path).unwrap();
74            f.read_to_string(&mut text)
75                .expect("Couldn't read source file");
76
77            // FIXME Rustfmt bug in string literals
78            #[rustfmt::skip]
79            let text = {
80                regex.replace_all(
81                    &text,
82                    "if $1 == ::protobuf::wire_format::WireTypeVarint {\
83                        $3 = $2.read_enum()?;\
84                    } else {\
85                        return ::std::result::Result::Err(::protobuf::rt::unexpected_wire_type(wire_type));\
86                    }",
87                )
88            };
89            let mut out = File::create(&path).unwrap();
90            out.write_all(text.as_bytes())
91                .expect("Could not write source file");
92        });
93    }
94
95    #[cfg(feature = "grpcio-protobuf-codec")]
96    fn import_grpcio(&self) {
97        use std::collections::BTreeMap;
98        use std::fs::OpenOptions;
99
100        if !self.re_export_services {
101            return;
102        }
103
104        // TODO should be behind an option
105        let paths: BTreeMap<_, _> = self
106            .list_rs_files()
107            .map(|path| (path.file_stem().unwrap().to_str().unwrap().to_owned(), path))
108            .collect();
109        for (name, path) in &paths {
110            if name.starts_with("wrapper_")
111                || *name == "mod"
112                || name.ends_with("_grpc")
113                || !paths.contains_key(&*format!("{}_grpc", name))
114            {
115                continue;
116            }
117
118            let mut out = OpenOptions::new()
119                .append(true)
120                .open(&path)
121                .expect("Couldn't open source file");
122            writeln!(out, "pub use super::{}_grpc::*;", name).expect("Could not write source file");
123        }
124    }
125
126    #[cfg(not(feature = "grpcio-protobuf-codec"))]
127    fn import_grpcio(&self) {}
128
129    #[cfg(feature = "grpcio-protobuf-codec")]
130    fn generate_grpcio(
131        &self,
132        desc: &[protobuf::descriptor::FileDescriptorProto],
133        files_to_generate: &[String],
134    ) {
135        let output_dir = std::path::Path::new(&self.out_dir);
136        let results = grpcio_compiler::codegen::gen(desc, &files_to_generate);
137        for res in results {
138            let out_file = output_dir.join(&res.name);
139            let mut f = std::fs::File::create(&out_file).unwrap();
140            f.write_all(&res.content).unwrap();
141        }
142    }
143
144    #[cfg(not(feature = "grpcio-protobuf-codec"))]
145    fn generate_grpcio(&self, _: &[protobuf::descriptor::FileDescriptorProto], _: &[String]) {}
146}
protobuf_build/protobuf_impl.rs

protobuf_build/
protobuf_impl.rs
