-
Notifications
You must be signed in to change notification settings - Fork 892
/
generate_check_code_prefix.rs
126 lines (116 loc) · 3.87 KB
/
generate_check_code_prefix.rs
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
//! Regenerate `src/checks_gen.rs` (namely, the `CheckCodePrefix` enum).
use std::collections::{BTreeMap, BTreeSet};
use std::fs::OpenOptions;
use std::io::{BufWriter, Write};
use anyhow::Result;
use codegen::{Scope, Type, Variant};
use itertools::Itertools;
use ruff::checks::CheckCode;
use strum::IntoEnumIterator;
fn main() -> Result<()> {
// Build up a map from prefix to matching CheckCodes.
let mut prefix_to_codes: BTreeMap<String, BTreeSet<CheckCode>> = Default::default();
for check_code in CheckCode::iter() {
let as_ref: String = check_code.as_ref().to_string();
let prefix_len = as_ref
.chars()
.take_while(|char| char.is_alphabetic())
.count();
for i in prefix_len..=as_ref.len() {
let prefix = as_ref[..i].to_string();
let entry = prefix_to_codes
.entry(prefix)
.or_insert_with(|| Default::default());
entry.insert(check_code.clone());
}
}
let mut scope = Scope::new();
// Create the `CheckCodePrefix` definition.
let mut gen = scope
.new_enum("CheckCodePrefix")
.vis("pub")
.derive("EnumString")
.derive("Debug")
.derive("PartialEq")
.derive("Eq")
.derive("PartialOrd")
.derive("Ord")
.derive("Clone")
.derive("Serialize")
.derive("Deserialize");
for (prefix, _) in &prefix_to_codes {
gen = gen.push_variant(Variant::new(prefix.to_string()));
}
// Create the `PrefixSpecificity` definition.
scope
.new_enum("PrefixSpecificity")
.vis("pub")
.derive("PartialEq")
.derive("Eq")
.derive("PartialOrd")
.derive("Ord")
.push_variant(Variant::new("Category"))
.push_variant(Variant::new("Hundreds"))
.push_variant(Variant::new("Tens"))
.push_variant(Variant::new("Explicit"));
// Create the `match` statement, to map from definition to relevant codes.
let mut gen = scope
.new_impl("CheckCodePrefix")
.new_fn("codes")
.arg_ref_self()
.ret(Type::new("Vec<CheckCode>"))
.vis("pub")
.line("match self {");
for (prefix, codes) in &prefix_to_codes {
gen = gen.line(format!(
"CheckCodePrefix::{prefix} => vec![{}],",
codes
.iter()
.map(|code| format!("CheckCode::{}", code.as_ref()))
.join(", ")
));
}
gen.line("}");
// Create the `match` statement, to map from definition to specificity.
let mut gen = scope
.new_impl("CheckCodePrefix")
.new_fn("specificity")
.arg_ref_self()
.ret(Type::new("PrefixSpecificity"))
.vis("pub")
.line("match self {");
for (prefix, _) in &prefix_to_codes {
let num_numeric = prefix.chars().filter(|char| char.is_numeric()).count();
let specificity = match num_numeric {
3 => "Explicit",
2 => "Tens",
1 => "Hundreds",
0 => "Category",
_ => panic!("Invalid prefix: {}", prefix),
};
gen = gen.line(format!(
"CheckCodePrefix::{prefix} => PrefixSpecificity::{},",
specificity
));
}
gen.line("}");
// Write the output to `src/checks_gen.rs`.
let f = OpenOptions::new()
.write(true)
.truncate(true)
.open("src/checks_gen.rs")
.expect("unable to open file");
let mut f = BufWriter::new(f);
writeln!(
f,
"//! File automatically generated by examples/generate_check_code_prefix.rs."
)?;
writeln!(f)?;
writeln!(f, "use serde::{{Serialize, Deserialize}};")?;
writeln!(f, "use strum_macros::EnumString;")?;
writeln!(f)?;
writeln!(f, "use crate::checks::CheckCode;")?;
writeln!(f)?;
writeln!(f, "{}", scope.to_string())?;
Ok(())
}