From d5fbcd708d1ae5760ffe90a7c0bfd94d6a9c6cc5 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Sat, 5 Nov 2022 14:30:18 -0400
Subject: [PATCH] Automatically write to src/checks_gen.rs (#604)

---
 examples/generate_check_code_prefix.rs | 54 ++++++++++++++++++++------
 1 file changed, 43 insertions(+), 11 deletions(-)

diff --git a/examples/generate_check_code_prefix.rs b/examples/generate_check_code_prefix.rs
index 5c97d1a0c5..bed8132148 100644
--- a/examples/generate_check_code_prefix.rs
+++ b/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(())
 }