Skip to content

Commit

Permalink
build: Add option to generate a full crate
Browse files Browse the repository at this point in the history
  • Loading branch information
Tuetuopay committed Jun 25, 2021
1 parent 3cc97f9 commit 12d0e93
Show file tree
Hide file tree
Showing 2 changed files with 107 additions and 8 deletions.
72 changes: 65 additions & 7 deletions prost-build/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ mod lib_generator;
mod message_graph;
mod path;

use std::collections::HashMap;
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::default;
use std::env;
use std::ffi::{OsStr, OsString};
Expand Down Expand Up @@ -228,7 +228,8 @@ const DISABLE_COMMENTS: &str = "disable_comments";
const EXTERN_PATH: &str = "extern_path";
const RETAIN_ENUM_PREFIX: &str = "retain_enum_prefix";
const MOD_RS: &str = "mod_rs";
pub const PROTOC_OPTS: [&str; 9] = [
const GEN_CRATE: &str = "gen_crate";
pub const PROTOC_OPTS: [&str; 10] = [
BTREE_MAP,
BYTES,
FIELD_ATTR,
Expand All @@ -238,6 +239,7 @@ pub const PROTOC_OPTS: [&str; 9] = [
EXTERN_PATH,
RETAIN_ENUM_PREFIX,
MOD_RS,
GEN_CRATE,
];

/// Configuration options for Protobuf code generation.
Expand All @@ -257,6 +259,7 @@ pub struct Config {
protoc_args: Vec<OsString>,
disable_comments: PathMap<()>,
mod_rs: bool,
manifest_tpl: Option<PathBuf>,
}

impl Config {
Expand Down Expand Up @@ -294,6 +297,7 @@ impl Config {
[EXTERN_PATH, k, v] => self.extern_paths.push((k.to_string(), v.to_string())),
[RETAIN_ENUM_PREFIX] => self.strip_enum_prefix = false,
[MOD_RS] => self.mod_rs = true,
[GEN_CRATE, v] => self.manifest_tpl = Some(v.into()),
_ if log_unknown => eprintln!("prost: Unknown option `{}`", opt.join("=")),
_ => (),
}
Expand Down Expand Up @@ -747,6 +751,20 @@ impl Config {
self
}

/// Generate a full crate based on a `Cargo.toml` template.
///
/// The output directory will contain the crate and not the built files directly, in `gen/`
/// folder. Each `proto` package will get its own feature flag, with proper dependency solving.
/// The `Cargo.toml` must contain `{{ features }}` string where all crate features will be
/// inserted.
pub fn generate_crate<P>(&mut self, manifest_template: P) -> &mut Self
where
P: Into<PathBuf>,
{
self.manifest_tpl = Some(manifest_template.into());
self
}

/// Compile `.proto` files into Rust files during a Cargo build with additional code generator
/// configuration options.
///
Expand Down Expand Up @@ -838,8 +856,7 @@ impl Config {

let modules = self.generate(file_descriptor_set.file)?;
for (module, content) in modules {
let mut filename = module.join(".");
filename.push_str(".rs");
let filename = self.filename(&module);

let output_path = target.join(&filename);

Expand All @@ -865,7 +882,7 @@ impl Config {
match self.generate(req.proto_file) {
Ok(modules) => {
let f = modules.into_iter().map(|(module, content)| File {
name: Some(module.join(".") + ".rs"),
name: Some(self.filename(&module)),
insertion_point: None,
content: Some(content),
generated_code_info: None,
Expand All @@ -884,9 +901,24 @@ impl Config {
}
}

fn filename(&self, module: &Module) -> String {
let (prefix, suffix) = match (self.manifest_tpl.is_some(), module.as_slice()) {
(true, [n]) if n == "lib" => ("src/", ".rs"),
(true, [n]) if n == "Cargo.toml" => ("", ""),
(true, _) => ("gen/", ".rs"),
(false, _) => ("", ".rs"),
};
prefix.to_owned() + &module.join(".") + suffix
}

fn generate(&mut self, files: Vec<FileDescriptorProto>) -> Result<HashMap<Module, String>> {
let mut modules = HashMap::new();
let mut packages = HashMap::new();
let deps = self
.manifest_tpl
.is_some()
.then(|| self.build_deps(&files))
.unwrap_or_default();

let message_graph = MessageGraph::new(&files)
.map_err(|error| Error::new(ErrorKind::InvalidInput, error))?;
Expand All @@ -912,14 +944,19 @@ impl Config {
}
}

if self.mod_rs {
if self.mod_rs || self.manifest_tpl.is_some() {
let mut mods = Mod::default();
for (module, _) in &modules {
mods.push(module);
}
let mut buf = modules.entry(vec!["mod".to_owned()]).or_default();
let name = self.manifest_tpl.is_some().then(|| "lib").unwrap_or("mod");
let mut buf = modules.entry(vec![name.to_owned()]).or_default();
LibGenerator::generate_librs(self, &mods, &mut buf);
}
if let Some(tpl) = self.manifest_tpl.clone() {
let mut buf = modules.entry(vec!["Cargo.toml".to_owned()]).or_default();
LibGenerator::generate_manifest(self, tpl, deps, &mut buf);
}

Ok(modules)
}
Expand All @@ -931,6 +968,26 @@ impl Config {
.map(to_snake)
.collect()
}

fn build_deps(&self, files: &[FileDescriptorProto]) -> BTreeMap<String, BTreeSet<String>> {
let names: HashMap<_, _> = files
.iter()
.map(|file| (file.name().to_owned(), file.package().to_owned()))
.collect();
let mut deps: BTreeMap<_, BTreeSet<_>> = BTreeMap::new();
for file in files {
let names = file
.dependency
.iter()
.filter_map(|dep| names.get(dep))
.filter(|dep| *dep != file.package())
.map(|s| s.clone());
deps.entry(file.package().to_owned())
.or_default()
.extend(names);
}
deps
}
}

impl default::Default for Config {
Expand All @@ -949,6 +1006,7 @@ impl default::Default for Config {
protoc_args: Vec::new(),
disable_comments: PathMap::default(),
mod_rs: false,
manifest_tpl: None,
}
}
}
Expand Down
43 changes: 42 additions & 1 deletion prost-build/src/lib_generator.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
use std::fs::read_to_string;
use std::path::PathBuf;

use crate::{Config, Module};

Expand Down Expand Up @@ -64,8 +66,19 @@ impl<'a> LibGenerator<'a> {
}

for package in mods.contents.iter().map(|content| content.join(".")) {
if self.config.manifest_tpl.is_some() {
let feature = package.replace("r#", "").replace(".", "_");
self.push_indent();
self.buf.push_str("#[cfg(feature = \"");
self.buf.push_str(&feature);
self.buf.push_str("\")]\n");
}

self.push_indent();
self.buf.push_str("include!(\"");
if self.config.manifest_tpl.is_some() {
self.buf.push_str("../gen/");
}
self.buf.push_str(&package);
self.buf.push_str(".rs\");\n");
}
Expand All @@ -76,4 +89,32 @@ impl<'a> LibGenerator<'a> {
self.buf.push_str(" ");
}
}

pub fn generate_manifest(
config: &'a mut Config,
template: PathBuf,
deps: BTreeMap<String, BTreeSet<String>>,
buf: &mut String,
) {
let mut generator = LibGenerator {
config,
depth: 0,
buf,
};
generator.push_manifest(template, deps);
}

fn push_manifest(&mut self, template: PathBuf, deps: BTreeMap<String, BTreeSet<String>>) {
let template = read_to_string(template).unwrap();
let mut buf = String::new();
for (feat, deps) in deps {
buf.push('"');
buf.push_str(&feat);
buf.push_str("\" = [");
let deps: Vec<_> = deps.iter().map(|dep| format!("\"{}\"", dep)).collect();
buf.push_str(&deps.join(", "));
buf.push_str("]\n");
}
self.buf.push_str(&template.replace("{{ features }}", &buf));
}
}

0 comments on commit 12d0e93

Please sign in to comment.