Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Added custom enum derive support #266

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Added custom enum derive support #266

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions pb-rs/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ pub struct ConfigBuilder {
error_cycle: bool,
headers: bool,
dont_use_cow: bool,
custom_enum_derive: Vec<String>,
custom_struct_derive: Vec<String>,
custom_repr: Option<String>,
owned: bool,
Expand Down Expand Up @@ -146,6 +147,12 @@ impl ConfigBuilder {
self
}

/// Add custom values to `#[derive(...)]` at the beginning of every enum
pub fn custom_enum_derive(mut self, val: Vec<String>) -> Self {
self.custom_enum_derive = val;
self
}

/// Add custom values to `#[derive(...)]` at the beginning of every structure
pub fn custom_struct_derive(mut self, val: Vec<String>) -> Self {
self.custom_struct_derive = val;
Expand Down Expand Up @@ -226,6 +233,7 @@ impl ConfigBuilder {
error_cycle: self.error_cycle,
headers: self.headers,
dont_use_cow: self.dont_use_cow, //Change this to true to not use cow with ./generate.sh for v2 and v3 tests
custom_enum_derive: self.custom_enum_derive.clone(),
custom_struct_derive: self.custom_struct_derive.clone(),
custom_repr: self.custom_repr.clone(),
custom_rpc_generator: Box::new(|_, _| Ok(())),
Expand Down
14 changes: 14 additions & 0 deletions pb-rs/src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,13 @@ fn run() -> Result<(), Error> {
.short("H")
.required(false)
.help("Do not add module comments and module attributes in generated file"),
).arg(
Arg::with_name("CUSTOM_ENUM_DERIVE")
.long("custom_enum_derive")
.short("E")
.required(false)
.takes_value(true)
.help("The comma separated values to add to #[derive(...)] for every enum"),
).arg(
Arg::with_name("CUSTOM_STRUCT_DERIVE")
.long("custom_struct_derive")
Expand Down Expand Up @@ -125,6 +132,12 @@ fn run() -> Result<(), Error> {
.split(',')
.map(|s| s.to_string())
.collect();
let custom_enum_derive: Vec<String> = matches
.value_of("CUSTOM_ENUM_DERIVE")
.unwrap_or("")
.split(',')
.map(|s| s.to_string())
.collect();

let compiler = ConfigBuilder::new(
&in_files,
Expand All @@ -138,6 +151,7 @@ fn run() -> Result<(), Error> {
.headers(!matches.is_present("NO_HEADERS"))
.dont_use_cow(matches.is_present("DONT_USE_COW"))
.custom_struct_derive(custom_struct_derive)
.custom_enum_derive(custom_enum_derive)
.nostd(matches.is_present("NOSTD"))
.hashbrown(matches.is_present("HASHBROWN"))
.gen_info(matches.is_present("GEN_INFO"))
Expand Down
55 changes: 33 additions & 22 deletions pb-rs/src/types.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1563,7 +1563,7 @@ impl Message {
m.write(w, desc, config)?;
}
for e in &self.enums {
e.write(w)?;
e.write(w, config)?;
}
for o in &self.oneofs {
o.write(w, desc, config)?;
Expand All @@ -1582,23 +1582,23 @@ impl Message {
desc: &FileDescriptor,
config: &Config,
) -> Result<()> {
let mut custom_struct_derive = config.custom_struct_derive.join(", ");

if !self.must_generate_impl_default(desc, config) {
custom_struct_derive += "Default";
}

if !custom_struct_derive.is_empty() {
custom_struct_derive += ", ";
}
let derives = if self.must_generate_impl_default(desc, config) {
vec!["Debug", "PartialEq", "Clone"]
} else {
vec!["Default", "Debug", "PartialEq", "Clone"]
};
let derives_str = config
.custom_struct_derive
.iter()
.map(|s| s.as_str())
.filter(|s| !s.is_empty())
.chain(derives.into_iter())
.collect::<Vec<_>>()
.join(", ");

writeln!(w, "#[allow(clippy::derive_partial_eq_without_eq)]")?;

writeln!(
w,
"#[derive({}Debug, PartialEq, Clone)]",
custom_struct_derive
)?;
writeln!(w, "#[derive({derives_str})]")?;

if let Some(repr) = &config.custom_repr {
writeln!(w, "#[repr({})]", repr)?;
Expand Down Expand Up @@ -2163,10 +2163,10 @@ impl Enumerator {
get_modules(&self.module, self.imported, desc)
}

fn write<W: Write>(&self, w: &mut W) -> Result<()> {
fn write<W: Write>(&self, w: &mut W, config: &Config) -> Result<()> {
println!("Writing enum {}", self.name);
writeln!(w)?;
self.write_definition(w)?;
self.write_definition(w, config)?;
writeln!(w)?;
if self.fields.is_empty() {
Ok(())
Expand All @@ -2179,8 +2179,18 @@ impl Enumerator {
}
}

fn write_definition<W: Write>(&self, w: &mut W) -> Result<()> {
writeln!(w, "#[derive(Debug, PartialEq, Eq, Clone, Copy)]")?;
fn write_definition<W: Write>(&self, w: &mut W, config: &Config) -> Result<()> {
let mut custom_enum_derive = config.custom_enum_derive.join(", ");

if !custom_enum_derive.is_empty() {
custom_enum_derive += ", ";
}

writeln!(
w,
"#[derive({}Debug, PartialEq, Eq, Clone, Copy)]",
custom_enum_derive
)?;
writeln!(w, "pub enum {} {{", self.name)?;
for (f, number) in &self.fields {
writeln!(w, " {} = {},", f, number)?;
Expand Down Expand Up @@ -2544,6 +2554,7 @@ pub struct Config {
pub error_cycle: bool,
pub headers: bool,
pub dont_use_cow: bool,
pub custom_enum_derive: Vec<String>,
pub custom_struct_derive: Vec<String>,
pub custom_repr: Option<String>,
pub custom_rpc_generator: RpcGeneratorFunction,
Expand Down Expand Up @@ -3030,7 +3041,7 @@ impl FileDescriptor {
self.write_package_start(w)?;
self.write_uses(w, config)?;
self.write_imports(w)?;
self.write_enums(w)?;
self.write_enums(w, config)?;
self.write_messages(w, config)?;
self.write_rpc_services(w, config)?;
self.write_package_end(w)?;
Expand Down Expand Up @@ -3116,11 +3127,11 @@ impl FileDescriptor {
Ok(())
}

fn write_enums<W: Write>(&self, w: &mut W) -> Result<()> {
fn write_enums<W: Write>(&self, w: &mut W, config: &Config) -> Result<()> {
for m in self.enums.iter().filter(|e| !e.imported) {
println!("Writing enum {}", m.name);
writeln!(w)?;
m.write_definition(w)?;
m.write_definition(w, config)?;
writeln!(w)?;
m.write_impl_default(w)?;
writeln!(w)?;
Expand Down
1 change: 1 addition & 0 deletions perftest/build.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ fn main() {
error_cycle: false,
headers: false,
dont_use_cow: false,
custom_enum_derive: vec![],
custom_struct_derive: vec![],
custom_repr: None,
custom_rpc_generator: Box::new(|rpc, writer| generate_rpc_test(rpc, writer)),
Expand Down