2023-09-24 16:27:06 +02:00
|
|
|
use anyhow::anyhow;
|
|
|
|
use genco::prelude::*;
|
|
|
|
use std::path::{Path, PathBuf};
|
|
|
|
use tokio::io::AsyncWriteExt;
|
|
|
|
use walkdir::WalkDir;
|
|
|
|
|
|
|
|
pub struct Codegen {}
|
|
|
|
|
|
|
|
impl Codegen {
|
|
|
|
pub fn new() -> Self {
|
|
|
|
Self {}
|
|
|
|
}
|
|
|
|
|
|
|
|
pub async fn generate_rust(&self, input_path: &Path, output_path: &Path) -> anyhow::Result<()> {
|
|
|
|
let input_protos = self.discover_files(input_path, "proto")?;
|
|
|
|
let (input_proto_paths, input_dir) = self.copy_protos(input_protos, input_path).await?;
|
|
|
|
let (output_proto_paths, output_dir) = self
|
|
|
|
.generate_rust_from_proto(input_proto_paths, output_path, input_dir.path())
|
|
|
|
.await?;
|
|
|
|
|
|
|
|
self.copy_rs(output_proto_paths, output_path, output_dir.path())
|
|
|
|
.await?;
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
|
|
|
fn discover_files(&self, input_path: &Path, extension: &str) -> anyhow::Result<Vec<PathBuf>> {
|
|
|
|
let mut proto_files = Vec::new();
|
|
|
|
for entry in WalkDir::new(input_path) {
|
|
|
|
let entry = entry?;
|
|
|
|
|
|
|
|
if let Some(extension) = entry.path().extension().and_then(|e| e.to_str()) {
|
|
|
|
proto_files.push(entry.into_path());
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
if proto_files.is_empty() {
|
|
|
|
anyhow::anyhow!(
|
|
|
|
"failed to find any protobuf files in: {}",
|
|
|
|
input_path.display()
|
|
|
|
);
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok(proto_files)
|
|
|
|
}
|
|
|
|
|
|
|
|
async fn copy_protos(
|
|
|
|
&self,
|
|
|
|
input_protos: Vec<PathBuf>,
|
|
|
|
root_path: &Path,
|
|
|
|
) -> anyhow::Result<(Vec<PathBuf>, tempfile::TempDir)> {
|
|
|
|
let in_tempdir = tempfile::TempDir::new()?;
|
|
|
|
let in_tempdir_path = in_tempdir.path();
|
|
|
|
|
|
|
|
let mut input_proto_paths = Vec::new();
|
|
|
|
for input_proto in &input_protos {
|
|
|
|
let rel_proto_path = input_proto.strip_prefix(root_path)?;
|
|
|
|
let in_proto_path = in_tempdir_path.join(rel_proto_path);
|
|
|
|
if let Some(dir) = in_proto_path.parent() {
|
|
|
|
if !dir.exists() {
|
|
|
|
tokio::fs::create_dir_all(dir).await?;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
tokio::fs::copy(input_proto, &in_proto_path).await?;
|
|
|
|
input_proto_paths.push(in_proto_path);
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok((input_proto_paths, in_tempdir))
|
|
|
|
}
|
|
|
|
|
|
|
|
async fn generate_rust_from_proto(
|
|
|
|
&self,
|
|
|
|
input_proto_paths: Vec<PathBuf>,
|
|
|
|
output_path: &Path,
|
|
|
|
in_root_path: &Path,
|
|
|
|
) -> anyhow::Result<(Vec<PathBuf>, tempfile::TempDir)> {
|
|
|
|
let out_tempdir = tempfile::TempDir::new()?;
|
|
|
|
let out_tempdir_path = out_tempdir.path();
|
|
|
|
let handle = tokio::task::spawn_blocking({
|
|
|
|
let out_tempdir_path = out_tempdir_path.to_path_buf();
|
|
|
|
let in_root_path = in_root_path.to_path_buf();
|
|
|
|
move || {
|
|
|
|
prost_build::Config::new()
|
|
|
|
.out_dir(out_tempdir_path)
|
|
|
|
.compile_protos(input_proto_paths.as_slice(), &[in_root_path])?;
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
});
|
|
|
|
|
|
|
|
let result: anyhow::Result<()> = handle.await?;
|
|
|
|
result?;
|
|
|
|
|
|
|
|
let mut output_paths = self.discover_files(&out_tempdir_path, "rs")?;
|
|
|
|
|
|
|
|
let mod_path = self
|
|
|
|
.generate_mod_file(&out_tempdir_path, &output_paths)
|
|
|
|
.await?;
|
|
|
|
output_paths.push(mod_path);
|
|
|
|
|
|
|
|
Ok((output_paths, out_tempdir))
|
|
|
|
}
|
|
|
|
|
|
|
|
async fn generate_mod_file(
|
|
|
|
&self,
|
|
|
|
output_tempdir_path: &Path,
|
|
|
|
output_paths: &[PathBuf],
|
|
|
|
) -> anyhow::Result<PathBuf> {
|
|
|
|
let mod_path = output_tempdir_path.join("mod.rs");
|
|
|
|
let mut mod_file = tokio::fs::File::create(&mod_path).await?;
|
|
|
|
|
|
|
|
let mut includes: Vec<genco::lang::rust::Tokens> = Vec::new();
|
|
|
|
for generated_file in output_paths {
|
|
|
|
if let Some(name) = generated_file.file_name() {
|
|
|
|
let mod_name = generated_file
|
|
|
|
.file_stem()
|
|
|
|
.unwrap()
|
|
|
|
.to_ascii_lowercase()
|
|
|
|
.to_string_lossy()
|
|
|
|
.replace(".", "_")
|
|
|
|
.replace("-", "_");
|
|
|
|
|
|
|
|
let file_name = name.to_str().unwrap();
|
|
|
|
|
|
|
|
includes.push(genco::quote! {
|
|
|
|
pub mod $(mod_name) {
|
|
|
|
include!($(quoted(file_name)));
|
|
|
|
}
|
|
|
|
});
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
let mod_tokens: genco::lang::rust::Tokens = genco::quote! {
|
|
|
|
$(for tokens in includes join($['\n']) => $tokens)
|
|
|
|
};
|
|
|
|
let mod_contents = mod_tokens.to_file_string()?;
|
|
|
|
mod_file.write_all(mod_contents.as_bytes()).await?;
|
|
|
|
|
|
|
|
Ok(mod_path)
|
|
|
|
}
|
|
|
|
|
|
|
|
async fn copy_rs(
|
|
|
|
&self,
|
|
|
|
output_proto_paths: Vec<PathBuf>,
|
|
|
|
output_path: &Path,
|
|
|
|
root_path: &Path,
|
|
|
|
) -> anyhow::Result<()> {
|
|
|
|
for output_rs in &output_proto_paths {
|
|
|
|
let rel_proto_path = output_rs.strip_prefix(root_path).map_err(|e| {
|
|
|
|
anyhow!(
|
|
|
|
"output: {} does not match root_path: {}",
|
|
|
|
output_rs.display(),
|
|
|
|
root_path.display()
|
|
|
|
)
|
|
|
|
})?;
|
|
|
|
let in_proto_path = output_path.join(rel_proto_path);
|
|
|
|
if let Some(dir) = in_proto_path.parent() {
|
|
|
|
if !dir.exists() {
|
|
|
|
tokio::fs::create_dir_all(dir).await?;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
tokio::fs::copy(output_rs, &in_proto_path).await?;
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl Default for Codegen {
|
|
|
|
fn default() -> Self {
|
|
|
|
Self::new()
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2023-09-24 14:02:01 +02:00
|
|
|
#[cfg(test)]
|
|
|
|
mod tests {
|
2023-09-24 16:27:06 +02:00
|
|
|
use genco::prelude::*;
|
|
|
|
use tokio::io::AsyncWriteExt;
|
|
|
|
|
2023-09-24 14:09:24 +02:00
|
|
|
#[tokio::test]
|
2023-09-24 16:27:06 +02:00
|
|
|
async fn test_can_generate_output_rust() -> anyhow::Result<()> {
|
|
|
|
// Generate from protobuf
|
2023-09-24 14:09:24 +02:00
|
|
|
let proto_spec = r#"
|
|
|
|
syntax = "proto3";
|
2023-09-24 14:02:01 +02:00
|
|
|
|
2023-09-24 16:27:06 +02:00
|
|
|
import "includes/test_include.proto";
|
|
|
|
|
2023-09-24 14:09:24 +02:00
|
|
|
package test.can.generate.output.rust;
|
|
|
|
|
|
|
|
message MyEvent {
|
|
|
|
string name = 1;
|
|
|
|
}
|
|
|
|
"#;
|
2023-09-24 16:27:06 +02:00
|
|
|
|
|
|
|
let proto_include_spec = r#"
|
|
|
|
syntax = "proto3";
|
|
|
|
|
|
|
|
package test.can.generate.output.rust.include.test_include;
|
|
|
|
|
|
|
|
message MyInclude {
|
|
|
|
string name = 1;
|
|
|
|
}
|
|
|
|
"#;
|
|
|
|
|
|
|
|
let out_tempdir = tempfile::TempDir::new()?;
|
|
|
|
let in_tempdir = tempfile::TempDir::new()?;
|
|
|
|
|
|
|
|
let proto_path = in_tempdir.path().join("test.proto");
|
|
|
|
let mut proto_file = tokio::fs::File::create(&proto_path).await?;
|
|
|
|
proto_file.write_all(proto_spec.as_bytes()).await?;
|
|
|
|
proto_file.sync_all().await?;
|
|
|
|
|
|
|
|
tokio::fs::create_dir_all(in_tempdir.path().join("includes")).await?;
|
|
|
|
let proto_include_path = in_tempdir.path().join("includes/test_include.proto");
|
|
|
|
let mut proto_file = tokio::fs::File::create(&proto_include_path).await?;
|
|
|
|
proto_file.write_all(proto_include_spec.as_bytes()).await?;
|
|
|
|
proto_file.sync_all().await?;
|
|
|
|
|
|
|
|
let out_tempdir_path = out_tempdir.into_path();
|
|
|
|
let handle = tokio::task::spawn_blocking({
|
|
|
|
let out_tempdir_path = out_tempdir_path.clone();
|
|
|
|
move || {
|
|
|
|
prost_build::Config::new()
|
|
|
|
.out_dir(out_tempdir_path)
|
|
|
|
.compile_protos(&[proto_path, proto_include_path], &[in_tempdir.into_path()])?;
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
});
|
|
|
|
|
|
|
|
let result: anyhow::Result<()> = handle.await?;
|
|
|
|
result?;
|
|
|
|
|
|
|
|
let mut entries = tokio::fs::read_dir(&out_tempdir_path).await?;
|
|
|
|
let mut file_paths = Vec::new();
|
|
|
|
while let Some(entry) = entries.next_entry().await? {
|
|
|
|
if let Some(ext) = entry.path().extension().and_then(|e| e.to_str()) {
|
|
|
|
if ext == "rs" {
|
|
|
|
file_paths.push(entry.path());
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
// Generate mod.rs
|
|
|
|
let mod_path = out_tempdir_path.join("mod.rs");
|
|
|
|
let mut mod_file = tokio::fs::File::create(&mod_path).await?;
|
|
|
|
|
|
|
|
let mut includes: Vec<genco::lang::rust::Tokens> = Vec::new();
|
|
|
|
for generated_file in &file_paths {
|
|
|
|
if let Some(name) = generated_file.file_name() {
|
|
|
|
let mod_name = generated_file
|
|
|
|
.file_stem()
|
|
|
|
.unwrap()
|
|
|
|
.to_ascii_lowercase()
|
|
|
|
.to_string_lossy()
|
|
|
|
.replace(".", "_")
|
|
|
|
.replace("-", "_");
|
|
|
|
|
|
|
|
let file_name = name.to_str().unwrap();
|
|
|
|
|
|
|
|
includes.push(genco::quote! {
|
|
|
|
pub mod $(mod_name) {
|
|
|
|
include!($(quoted(file_name)));
|
|
|
|
}
|
|
|
|
});
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
let mod_tokens: genco::lang::rust::Tokens = genco::quote! {
|
|
|
|
$(for tokens in includes join($['\n']) => $tokens)
|
|
|
|
};
|
|
|
|
let mod_contents = mod_tokens.to_file_string()?;
|
|
|
|
|
|
|
|
pretty_assertions::assert_eq!("", mod_contents);
|
|
|
|
|
|
|
|
mod_file.write_all(mod_contents.as_bytes()).await?;
|
|
|
|
|
|
|
|
assert_eq!(1, file_paths.len());
|
|
|
|
|
|
|
|
Ok(())
|
2023-09-24 14:02:01 +02:00
|
|
|
}
|
|
|
|
}
|