Automatically write to src/checks_gen.rs (#604)

examples/generate_check_code_prefix.rs

@@ -1,13 +1,28 @@
-//! Generate the CheckCodePrefix enum.
-
 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;
 
-fn main() {
+#[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() {
@@ -100,12 +115,29 @@ fn main() {
     }
     gen.line("}");
 
-    println!("//! File automatically generated by examples/generate_check_code_prefix.rs.");
-    println!();
-    println!("use serde::{{Serialize, Deserialize}};");
-    println!("use strum_macros::EnumString;");
-    println!();
-    println!("use crate::checks::CheckCode;");
-    println!();
-    println!("{}", scope.to_string());
+    // 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(())
 }