grpc_build/
base.rs

1use std::{
2    ffi::OsString,
3    path::{Path, PathBuf},
4    process::Command,
5};
6
7use anyhow::{Context, Result};
8use walkdir::WalkDir;
9
10// the use of these `inner` functions is a compile time optimisation. In this case it's probably
11// minimal but it improves how the code compiles. The inner functions are not generic, so can be built exactly once
12// but the outer functions are generic and must be built for every input type (String, &String, &str, &Path, etc).
13// Since the outer function just calls the inner function, this is very cheap, but still provides the ergonomic generic API
14
15pub fn prepare_out_dir(out_dir: impl AsRef<Path>) -> Result<()> {
16    fn inner(out_dir: &Path) -> Result<()> {
17        if out_dir.exists() {
18            fs_err::remove_dir_all(out_dir).with_context(|| {
19                format!(
20                    "could not remove the output directory: {}",
21                    out_dir.display()
22                )
23            })?;
24        }
25
26        fs_err::create_dir_all(out_dir).with_context(|| {
27            format!(
28                "could not create the output directory: {}",
29                out_dir.display()
30            )
31        })?;
32
33        Ok(())
34    }
35    inner(out_dir.as_ref())
36}
37
38/// Get all the `.proto` files within the provided directory
39#[allow(clippy::filetype_is_file)]
40pub fn get_protos(input: impl AsRef<Path>, follow_links: bool) -> impl Iterator<Item = PathBuf> {
41    fn inner(input: &Path, follow_links: bool) -> impl Iterator<Item = PathBuf> {
42        // TODO: maybe add this?
43        // println!("cargo:rerun-if-changed={}", input.display());
44
45        WalkDir::new(input)
46            .follow_links(follow_links)
47            .into_iter()
48            .filter_map(|r| r.map_err(|err| println!("cargo:warning={err:?}")).ok())
49            .filter(|e| e.file_type().is_file())
50            .filter(|e| e.path().extension().is_some_and(|e| e == "proto"))
51            .map(|e| e.path().to_path_buf())
52    }
53    inner(input.as_ref(), follow_links)
54}
55
56/// [`tonic_build::Builder::compile`] outputs all the rust files into the output dir all at the top level.
57/// This might not be the most desirable. Running this function converts the file into a more expected directory
58/// structure and generates the expected mod file output
59pub fn refactor(output: impl AsRef<Path>) -> Result<()> {
60    fn inner(output: &Path) -> Result<()> {
61        let tree: crate::tree::Tree = fs_err::read_dir(output)?
62            .filter_map(|r| r.map_err(|err| println!("cargo:warning={err:?}")).ok())
63            .filter(|e| e.path().extension().is_some_and(|e| e == "rs"))
64            .filter(|e| !e.path().ends_with("mod.rs"))
65            .map(|e| e.path())
66            .collect();
67
68        tree.move_paths(output, OsString::new(), PathBuf::new())?;
69        fs_err::write(output.join("mod.rs"), tree.generate_module())?;
70
71        Command::new("rustfmt")
72            .arg(output.join("mod.rs"))
73            .spawn()
74            .context("failed to format the mod.rs output")?;
75
76        Ok(())
77    }
78    inner(output.as_ref())
79}
80
81#[cfg(test)]
82mod test {
83    use super::refactor;
84
85    #[test]
86    fn refactor_test_moves_files_to_correct_place() {
87        let files = vec![
88            "root.pak.a1.rs",
89            "root.pak.a2.rs",
90            "root.pak.rs",
91            "root.now.deeply.nested.rs",
92            "root.rs",
93            "other.rs",
94        ];
95
96        let temp_dir = tempfile::tempdir().unwrap();
97        let temp_dir_path = temp_dir.path().to_path_buf();
98        // create files
99        for file in &files {
100            let path = temp_dir_path.join(file);
101            std::fs::create_dir_all(path.parent().unwrap()).unwrap();
102            std::fs::File::create(path.clone()).unwrap();
103            // write file name as its content
104            std::fs::write(path.clone(), format!("// {file} contents")).unwrap();
105        }
106
107        let expected_file_contents = vec![
108            ("root/pak/a1.rs", vec!["// root.pak.a1.rs contents"]),
109            ("root/pak/a2.rs", vec!["// root.pak.a2.rs content"]),
110            (
111                "root/pak.rs",
112                vec!["pub mod a1;", "pub mod a2;", "// root.pak.rs contents"],
113            ),
114            ("root/now.rs", vec!["pub mod deeply;"]),
115            ("root/now/deeply.rs", vec!["pub mod nested;"]),
116            (
117                "root/now/deeply/nested.rs",
118                vec!["// root.now.deeply.nested.rs contents"],
119            ),
120            (
121                "root.rs",
122                vec!["pub mod pak;", "pub mod now;", "// root.rs contents"],
123            ),
124            ("mod.rs", vec!["pub mod other;", "pub mod root;"]),
125            ("other.rs", vec!["// other.rs contents"]),
126        ];
127
128        // Act
129        refactor(&temp_dir_path).unwrap();
130
131        // check if files are moved and contents are correct
132        for (file, contents) in &expected_file_contents {
133            let path = temp_dir_path.join(file);
134            assert!(path.exists());
135            let content = std::fs::read_to_string(path).unwrap();
136            for line in contents {
137                assert!(content.contains(line), "{content} does not contain {line}");
138            }
139        }
140    }
141}