ruff/examples/generate_check_code_prefix.rs

use std::collections::{BTreeMap, BTreeSet};
use std::fs::OpenOptions;
use std::io;
use std::io::Write;

use anyhow::Result;
use clap::Parser;
use codegen::{Scope, Type, Variant};
use itertools::Itertools;
use ruff::checks::CheckCode;
use strum::IntoEnumIterator;

#[derive(Parser)]
#[command(author, version, about, long_about = None)]
/// Generate the `CheckCodePrefix` enum.
struct Cli {
    /// Write the generated source code to stdout (rather than to
    /// `src/checks_gen.rs`).
    #[arg(long)]
    dry_run: bool,
}

fn main() -> Result<()> {
    let cli = Cli::parse();

    // 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 mut writer = if cli.dry_run {
        Box::new(io::stdout()) as Box<dyn Write>
    } else {
        let f = OpenOptions::new()
            .write(true)
            .truncate(true)
            .open("src/checks_gen.rs")
            .expect("unable to open file");
        Box::new(f) as Box<dyn Write>
    };

    writeln!(
        writer,
        "//! File automatically generated by examples/generate_check_code_prefix.rs."
    )?;
    writeln!(writer)?;
    writeln!(writer, "use serde::{{Serialize, Deserialize}};")?;
    writeln!(writer, "use strum_macros::EnumString;")?;
    writeln!(writer)?;
    writeln!(writer, "use crate::checks::CheckCode;")?;
    writeln!(writer)?;
    writeln!(writer, "{}", scope.to_string())?;

    Ok(())
}