From 0b76fdf5d0ff3f25371188c51807175c607f9e30 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Sat, 5 Nov 2022 14:23:13 -0400
Subject: [PATCH] Automatically write to src/checks_gen.rs

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

diff --git a/examples/generate_check_code_prefix.rs b/examples/generate_check_code_prefix.rs
index 5c97d1a0c5e598..1a338605e98cfd 100644
--- a/examples/generate_check_code_prefix.rs
+++ b/examples/generate_check_code_prefix.rs
@@ -1,13 +1,15 @@
-//! Generate the CheckCodePrefix enum.
+//! Regenerate `src/checks_gen.rs` (namely, the `CheckCodePrefix` enum).
 
 use std::collections::{BTreeMap, BTreeSet};
+use std::fs::{File, 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() {
+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() {
@@ -100,12 +102,24 @@ 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 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(())
 }