1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

//! A prost toolkit to build protobuf with serde support.
//!
//! Usually when we define our protobuf messages, we hope some of the generated
//! data structures have good serde support. Fortunately `serde-build` has that
//! capability - you can create a config with `prost_build::Config::new()` and
//! and then set proper attributes for type or field. For exmaple, you can add
//! serde support for this structure by using:
//!
//! ```ignore
//! config.type_attribute("package.RequestPing", "#[derive(serde::Serialize, serde::Deserialize)]");
//! config.type_attribute("package.RequestPing", "#[serde(default)]");
//! ```
//!
//! and you will get this generated code:
//!
//! ```ignore
//! #[derive(serde::Serialize, serde::Deserialize)]
//! #[serde(default)]
//! #[derive(Clone, PartialEq, ::prost::Message)]
//! pub struct RequestPing {
//!     #[prost(string, tag = "1")]
//!     pub ping: ::prost::alloc::string::String,
//! }
//! ```
//!
//! This crate helps to simplify this build script by using a predefined build configuration.
//!
//! # Getting started
//!
//! First of all, you shall create a JSON file which contains configuration for the builder. You can
//! get a copy of a default JSON file from: [default_build_config.json](https://raw.githubusercontent.com/tyrchen/prost-helper/master/prost-serde/default_build_config.json). See an example of [build_config.json](https://raw.githubusercontent.com/tyrchen/prost-helper/master/prost-serde/examples/build_config.json).
//! Please add the proto files, proto includes, output dir, and the data structure or field you want
//! to add the desired attributes.
//! Then you could use it in:
//!
//! ```ignore
//! use prost_serde::build_with_serde;
//!
//! fn main() {
//!     let json = include_str!("../examples/build_config.json");
//!     build_with_serde(json);
//! }
//! ```

use serde::{Deserialize, Serialize};
use std::{fs, process::Command};

#[derive(Deserialize, Serialize, Debug, Default)]
#[serde(default)]
pub struct BuildConfig {
    /// protobuf include dirs
    pub includes: Vec<String>,
    /// protobuf files
    pub files: Vec<String>,
    /// dir for generated code, defaults to Cargo OUT_DIR, else the current dir
    pub output: Option<String>,
    /// build options for serde support
    pub opts: Vec<BuildOption>,
}

#[derive(Deserialize, Serialize, Debug, Default)]
#[serde(default)]
pub struct BuildOption {
    /// scope of the attribute
    pub scope: String,
    /// description of the option
    pub description: String,
    /// extra attribute to put on generated data structure, for example: `#[derive(Serialize, Deserialize)]`
    pub attr: String,
    /// a list of paths you want to add the attribute
    pub paths: Vec<String>,
}

/// Build the protobuf files with the build opts provided by a JSON string.
///
/// Normally you should put the json file in your crate, and load it with `include_str!`,
/// then pass it to this build function.
pub fn build_with_serde(json: &str) -> BuildConfig {
    let build_config: BuildConfig = serde_json::from_str(json).unwrap();

    // For the output directory, use the specified one, or fallback to the
    // OUT_DIR env variable provided by Cargo if it exists (it should!), else
    // fallback to the current directory.
    let output_dir: String = match &build_config.output {
        None => {
            let default_output_dir = std::env::var("OUT_DIR");

            match default_output_dir {
                Err(_) => String::new(),
                Ok(cargo_out_dir) => cargo_out_dir,
            }
        }
        Some(specified_output) => specified_output.to_owned(),
    };

    let mut config = prost_build::Config::new();
    for opt in build_config.opts.iter() {
        match opt.scope.as_ref() {
            "bytes" => {
                config.bytes(&opt.paths);
                continue;
            }
            "btree_map" => {
                config.btree_map(&opt.paths);
                continue;
            }
            _ => (),
        };
        for path in opt.paths.iter() {
            match opt.scope.as_str() {
                "type" => config.type_attribute(path, opt.attr.as_str()),
                "field" => config.field_attribute(path, opt.attr.as_str()),
                v => panic!("Not supported type: {}", v),
            };
        }
    }

    fs::create_dir_all(&output_dir).unwrap();
    config.out_dir(&output_dir);

    config
        .compile_protos(&build_config.files, &build_config.includes)
        .unwrap_or_else(|e| panic!("Failed to compile proto files. Error: {:?}", e));

    Command::new("cargo")
        .args(&["fmt"])
        .status()
        .expect("cargo fmt failed");

    build_config
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn generate_serde_supported_code() {
        let json = include_str!("../examples/build_config.json");
        build_with_serde(json);
    }
}