#![doc(html_root_url = "https://docs.rs/ntex-prost-build/0.11.0")]
#![allow(clippy::option_as_ref_deref)]
mod ast;
mod code_generator;
mod extern_paths;
mod ident;
mod path;
use std::ffi::{OsStr, OsString};
use std::io::{Error, ErrorKind, Result, Write};
use std::ops::RangeToInclusive;
use std::path::{Path, PathBuf};
use std::{collections::HashMap, default, env, fmt, fs, process::Command};
use log::trace;
use prost::Message;
use prost_types::{FileDescriptorProto, FileDescriptorSet};
pub use crate::ast::{Comments, Method, Service};
use crate::code_generator::CodeGenerator;
use crate::extern_paths::ExternPaths;
use crate::ident::to_snake;
use crate::path::PathMap;
pub trait ServiceGenerator {
fn generate(&mut self, service: Service, buf: &mut String, priv_buf: &mut String);
fn finalize(&mut self, _buf: &mut String) {}
fn finalize_package(&mut self, _package: &str, _buf: &mut String) {}
}
pub struct Config {
file_descriptor_set_path: Option<PathBuf>,
service_generator: Option<Box<dyn ServiceGenerator>>,
types_map: PathMap<String>,
type_attributes: PathMap<String>,
field_attributes: PathMap<String>,
prost_types: bool,
strip_enum_prefix: bool,
out_dir: Option<PathBuf>,
extern_paths: Vec<(String, String)>,
default_package_filename: String,
protoc_args: Vec<OsString>,
disable_comments: PathMap<()>,
skip_protoc_run: bool,
include_file: Option<PathBuf>,
}
impl Config {
pub fn new() -> Config {
Config::default()
}
pub fn map_field_type<I, S>(&mut self, paths: I, tp: &str) -> &mut Self
where
I: IntoIterator<Item = S>,
S: AsRef<str>,
{
for matcher in paths {
self.types_map
.insert(matcher.as_ref().to_string(), tp.to_string());
}
self
}
pub fn field_attribute<P, A>(&mut self, path: P, attribute: A) -> &mut Self
where
P: AsRef<str>,
A: AsRef<str>,
{
self.field_attributes
.insert(path.as_ref().to_string(), attribute.as_ref().to_string());
self
}
pub fn type_attribute<P, A>(&mut self, path: P, attribute: A) -> &mut Self
where
P: AsRef<str>,
A: AsRef<str>,
{
self.type_attributes
.insert(path.as_ref().to_string(), attribute.as_ref().to_string());
self
}
pub fn service_generator(
&mut self,
service_generator: Box<dyn ServiceGenerator>,
) -> &mut Self {
self.service_generator = Some(service_generator);
self
}
pub fn compile_well_known_types(&mut self) -> &mut Self {
self.prost_types = false;
self
}
pub fn disable_comments<I, S>(&mut self, paths: I) -> &mut Self
where
I: IntoIterator<Item = S>,
S: AsRef<str>,
{
self.disable_comments.clear();
for matcher in paths {
self.disable_comments
.insert(matcher.as_ref().to_string(), ());
}
self
}
pub fn extern_path<P1, P2>(&mut self, proto_path: P1, rust_path: P2) -> &mut Self
where
P1: Into<String>,
P2: Into<String>,
{
self.extern_paths
.push((proto_path.into(), rust_path.into()));
self
}
pub fn file_descriptor_set_path<P>(&mut self, path: P) -> &mut Self
where
P: Into<PathBuf>,
{
self.file_descriptor_set_path = Some(path.into());
self
}
pub fn skip_protoc_run(&mut self) -> &mut Self {
self.skip_protoc_run = true;
self
}
pub fn retain_enum_prefix(&mut self) -> &mut Self {
self.strip_enum_prefix = false;
self
}
pub fn out_dir<P>(&mut self, path: P) -> &mut Self
where
P: Into<PathBuf>,
{
self.out_dir = Some(path.into());
self
}
pub fn default_package_filename<S>(&mut self, filename: S) -> &mut Self
where
S: Into<String>,
{
self.default_package_filename = filename.into();
self
}
pub fn protoc_arg<S>(&mut self, arg: S) -> &mut Self
where
S: AsRef<OsStr>,
{
self.protoc_args.push(arg.as_ref().to_owned());
self
}
pub fn include_file<P>(&mut self, path: P) -> &mut Self
where
P: Into<PathBuf>,
{
self.include_file = Some(path.into());
self
}
pub fn compile_protos(
&mut self,
protos: &[impl AsRef<Path>],
includes: &[impl AsRef<Path>],
) -> Result<()> {
let mut target_is_env = false;
let target: PathBuf = self.out_dir.clone().map(Ok).unwrap_or_else(|| {
env::var_os("OUT_DIR")
.ok_or_else(|| {
Error::new(ErrorKind::Other, "OUT_DIR environment variable is not set")
})
.map(|val| {
target_is_env = true;
Into::into(val)
})
})?;
let tmp;
let file_descriptor_set_path = if let Some(path) = &self.file_descriptor_set_path {
path.clone()
} else {
if self.skip_protoc_run {
return Err(Error::new(
ErrorKind::Other,
"file_descriptor_set_path is required with skip_protoc_run",
));
}
tmp = tempfile::Builder::new().prefix("prost-build").tempdir()?;
tmp.path().join("prost-descriptor-set")
};
if !self.skip_protoc_run {
let mut cmd = Command::new(protoc());
cmd.arg("--include_imports")
.arg("--include_source_info")
.arg("-o")
.arg(&file_descriptor_set_path);
for include in includes {
cmd.arg("-I").arg(include.as_ref());
}
cmd.arg("-I").arg(protoc_include());
for arg in &self.protoc_args {
cmd.arg(arg);
}
for proto in protos {
cmd.arg(proto.as_ref());
}
let output = cmd.output().map_err(|error| {
Error::new(
error.kind(),
format!("failed to invoke protoc (hint: https://docs.rs/prost-build/#sourcing-protoc): {}", error),
)
})?;
if !output.status.success() {
return Err(Error::new(
ErrorKind::Other,
format!("protoc failed: {}", String::from_utf8_lossy(&output.stderr)),
));
}
}
let buf = fs::read(file_descriptor_set_path)?;
let file_descriptor_set = FileDescriptorSet::decode(&*buf).map_err(|error| {
Error::new(
ErrorKind::InvalidInput,
format!("invalid FileDescriptorSet: {}", error),
)
})?;
let requests = file_descriptor_set
.file
.into_iter()
.map(|descriptor| {
(
Module::from_protobuf_package_name(descriptor.package()),
descriptor,
)
})
.collect::<Vec<_>>();
let file_names = requests
.iter()
.map(|req| {
(
req.0.clone(),
req.0.to_file_name_or(&self.default_package_filename),
)
})
.collect::<HashMap<Module, String>>();
let modules = self.generate(requests)?;
for (module, content) in &modules {
let file_name = file_names
.get(module)
.expect("every module should have a filename");
if file_name == "google.protobuf.rs" {
continue;
}
let output_path = target.join(file_name);
let previous_content = fs::read(&output_path);
if previous_content
.map(|previous_content| previous_content == content.as_bytes())
.unwrap_or(false)
{
trace!("unchanged: {:?}", file_name);
} else {
trace!("writing: {:?}", file_name);
fs::write(output_path, content)?;
}
}
if let Some(ref include_file) = self.include_file {
trace!("Writing include file: {:?}", target.join(include_file));
let mut file = fs::File::create(target.join(include_file))?;
self.write_includes(
modules.keys().collect(),
&mut file,
0,
if target_is_env { None } else { Some(&target) },
)?;
file.flush()?;
}
Ok(())
}
fn write_includes(
&self,
mut entries: Vec<&Module>,
outfile: &mut fs::File,
depth: usize,
basepath: Option<&PathBuf>,
) -> Result<usize> {
let mut written = 0;
while !entries.is_empty() {
let modident = entries[0].part(depth);
let matching: Vec<&Module> = entries
.iter()
.filter(|&v| v.part(depth) == modident)
.copied()
.collect();
{
let _temp = entries
.drain(..)
.filter(|&v| v.part(depth) != modident)
.collect();
entries = _temp;
}
self.write_line(outfile, depth, &format!("pub mod {} {{", modident))?;
let subwritten = self.write_includes(
matching
.iter()
.filter(|v| v.len() > depth + 1)
.copied()
.collect(),
outfile,
depth + 1,
basepath,
)?;
written += subwritten;
if subwritten != matching.len() {
let modname = matching[0].to_partial_file_name(..=depth);
if basepath.is_some() {
self.write_line(
outfile,
depth + 1,
&format!("include!(\"{}.rs\");", modname),
)?;
} else {
self.write_line(
outfile,
depth + 1,
&format!("include!(concat!(env!(\"OUT_DIR\"), \"/{}.rs\"));", modname),
)?;
}
written += 1;
}
self.write_line(outfile, depth, "}")?;
}
Ok(written)
}
fn write_line(&self, outfile: &mut fs::File, depth: usize, line: &str) -> Result<()> {
outfile.write_all(format!("{}{}\n", (" ").to_owned().repeat(depth), line).as_bytes())
}
pub fn generate(
&mut self,
requests: Vec<(Module, FileDescriptorProto)>,
) -> Result<HashMap<Module, String>> {
let mut modules = HashMap::new();
let mut packages = HashMap::new();
let extern_paths = ExternPaths::new(&self.extern_paths, self.prost_types)
.map_err(|error| Error::new(ErrorKind::InvalidInput, error))?;
for request in requests {
if !request.1.service.is_empty() {
packages.insert(request.0.clone(), request.1.package().to_string());
}
let buf = modules.entry(request.0).or_insert_with(String::new);
buf.insert_str(
0,
"#![allow(dead_code, unused_mut, unused_variables, clippy::identity_op, clippy::derivable_impls, clippy::unit_arg, clippy::derive_partial_eq_without_eq)]\n/// DO NOT MODIFY. Auto-generated file\n\n",
);
CodeGenerator::generate(self, &extern_paths, request.1, buf);
}
if let Some(ref mut service_generator) = self.service_generator {
for (module, package) in packages {
let buf = modules.get_mut(&module).unwrap();
service_generator.finalize_package(&package, buf);
}
}
Ok(modules)
}
}
impl default::Default for Config {
fn default() -> Config {
Config {
file_descriptor_set_path: None,
service_generator: None,
types_map: PathMap::default(),
type_attributes: PathMap::default(),
field_attributes: PathMap::default(),
prost_types: true,
strip_enum_prefix: true,
out_dir: None,
extern_paths: Vec::new(),
default_package_filename: "_".to_string(),
protoc_args: Vec::new(),
disable_comments: PathMap::default(),
skip_protoc_run: false,
include_file: None,
}
}
}
impl fmt::Debug for Config {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Config")
.field("file_descriptor_set_path", &self.file_descriptor_set_path)
.field("service_generator", &self.service_generator.is_some())
.field("types_map", &self.types_map)
.field("type_attributes", &self.type_attributes)
.field("field_attributes", &self.field_attributes)
.field("prost_types", &self.prost_types)
.field("strip_enum_prefix", &self.strip_enum_prefix)
.field("out_dir", &self.out_dir)
.field("extern_paths", &self.extern_paths)
.field("default_package_filename", &self.default_package_filename)
.field("protoc_args", &self.protoc_args)
.field("disable_comments", &self.disable_comments)
.finish()
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct Module {
components: Vec<String>,
}
impl Module {
pub fn from_parts<I>(parts: I) -> Self
where
I: IntoIterator,
I::Item: Into<String>,
{
Self {
components: parts.into_iter().map(|s| s.into()).collect(),
}
}
pub fn from_protobuf_package_name(name: &str) -> Self {
Self {
components: name
.split('.')
.filter(|s| !s.is_empty())
.map(to_snake)
.collect(),
}
}
pub fn parts(&self) -> impl Iterator<Item = &str> {
self.components.iter().map(|s| s.as_str())
}
pub fn to_file_name_or(&self, default: &str) -> String {
let mut root = if self.components.is_empty() {
default.to_owned()
} else {
self.components.join(".")
};
if !root.ends_with(".rs") {
root.push_str(".rs");
}
root
}
pub fn len(&self) -> usize {
self.components.len()
}
pub fn is_empty(&self) -> bool {
self.components.is_empty()
}
fn to_partial_file_name(&self, range: RangeToInclusive<usize>) -> String {
self.components[range].join(".")
}
fn part(&self, idx: usize) -> &str {
self.components[idx].as_str()
}
}
impl fmt::Display for Module {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let mut parts = self.parts();
if let Some(first) = parts.next() {
f.write_str(first)?;
}
for part in parts {
f.write_str("::")?;
f.write_str(part)?;
}
Ok(())
}
}
pub fn compile_protos(protos: &[impl AsRef<Path>], includes: &[impl AsRef<Path>]) -> Result<()> {
Config::new().compile_protos(protos, includes)
}
pub fn protoc() -> PathBuf {
match env::var_os("PROTOC") {
Some(protoc) => PathBuf::from(protoc),
None => PathBuf::from(env!("PROTOC")),
}
}
pub fn protoc_include() -> PathBuf {
match env::var_os("PROTOC_INCLUDE") {
Some(include) => PathBuf::from(include),
None => PathBuf::from(env!("PROTOC_INCLUDE")),
}
}