From 5694f3dc01d0a3af11eaa333526f4c854829c488 Mon Sep 17 00:00:00 2001
From: Clement Rey <cr.rey.clement@gmail.com>
Date: Wed, 18 Oct 2023 16:31:56 +0200
Subject: [PATCH 01/12] new CodeFormatter trait

---
 crates/re_types_builder/src/codegen/mod.rs   |  4 +---
 crates/re_types_builder/src/format/cpp.rs    | 11 +++++++++++
 crates/re_types_builder/src/format/mod.rs    | 15 +++++++++++++++
 crates/re_types_builder/src/format/python.rs | 11 +++++++++++
 crates/re_types_builder/src/format/rust.rs   | 11 +++++++++++
 crates/re_types_builder/src/lib.rs           | 13 +++++++++++--
 6 files changed, 60 insertions(+), 5 deletions(-)
 create mode 100644 crates/re_types_builder/src/format/cpp.rs
 create mode 100644 crates/re_types_builder/src/format/mod.rs
 create mode 100644 crates/re_types_builder/src/format/python.rs
 create mode 100644 crates/re_types_builder/src/format/rust.rs

diff --git a/crates/re_types_builder/src/codegen/mod.rs b/crates/re_types_builder/src/codegen/mod.rs
index 59354aaea158..2e405dfaad9b 100644
--- a/crates/re_types_builder/src/codegen/mod.rs
+++ b/crates/re_types_builder/src/codegen/mod.rs
@@ -1,5 +1,3 @@
-pub type GeneratedFiles = std::collections::BTreeMap<camino::Utf8PathBuf, String>;
-
 /// Implements the codegen pass.
 pub trait CodeGenerator {
     /// Generates user-facing code from [`crate::Objects`].
@@ -12,7 +10,7 @@ pub trait CodeGenerator {
         reporter: &crate::Reporter,
         objects: &crate::Objects,
         arrow_registry: &crate::ArrowRegistry,
-    ) -> GeneratedFiles;
+    ) -> crate::GeneratedFiles;
 }
 
 // ---
diff --git a/crates/re_types_builder/src/format/cpp.rs b/crates/re_types_builder/src/format/cpp.rs
new file mode 100644
index 000000000000..e282f7734047
--- /dev/null
+++ b/crates/re_types_builder/src/format/cpp.rs
@@ -0,0 +1,11 @@
+use crate::CodeFormatter;
+
+// ---
+
+pub struct CppCodeFormatter;
+
+impl CodeFormatter for CppCodeFormatter {
+    fn format(&mut self, reporter: &crate::Reporter, files: &mut crate::GeneratedFiles) {
+        todo!()
+    }
+}
diff --git a/crates/re_types_builder/src/format/mod.rs b/crates/re_types_builder/src/format/mod.rs
new file mode 100644
index 000000000000..ab274aeff520
--- /dev/null
+++ b/crates/re_types_builder/src/format/mod.rs
@@ -0,0 +1,15 @@
+/// Implements the formatting pass.
+pub trait CodeFormatter {
+    /// Formats generated files in-place.
+    fn format(&mut self, reporter: &crate::Reporter, files: &mut crate::GeneratedFiles);
+}
+
+// ---
+
+mod cpp;
+mod python;
+mod rust;
+
+pub use self::cpp::CppCodeFormatter;
+pub use self::python::PythonCodeFormatter;
+pub use self::rust::RustCodeFormatter;
diff --git a/crates/re_types_builder/src/format/python.rs b/crates/re_types_builder/src/format/python.rs
new file mode 100644
index 000000000000..c056f89ec4df
--- /dev/null
+++ b/crates/re_types_builder/src/format/python.rs
@@ -0,0 +1,11 @@
+use crate::CodeFormatter;
+
+// ---
+
+pub struct PythonCodeFormatter;
+
+impl CodeFormatter for PythonCodeFormatter {
+    fn format(&mut self, reporter: &crate::Reporter, files: &mut crate::GeneratedFiles) {
+        todo!()
+    }
+}
diff --git a/crates/re_types_builder/src/format/rust.rs b/crates/re_types_builder/src/format/rust.rs
new file mode 100644
index 000000000000..bf0c3b7b7bb7
--- /dev/null
+++ b/crates/re_types_builder/src/format/rust.rs
@@ -0,0 +1,11 @@
+use crate::CodeFormatter;
+
+// ---
+
+pub struct RustCodeFormatter;
+
+impl CodeFormatter for RustCodeFormatter {
+    fn format(&mut self, reporter: &crate::Reporter, files: &mut crate::GeneratedFiles) {
+        todo!()
+    }
+}
diff --git a/crates/re_types_builder/src/lib.rs b/crates/re_types_builder/src/lib.rs
index 68a7b71eba87..be6dfae8345b 100644
--- a/crates/re_types_builder/src/lib.rs
+++ b/crates/re_types_builder/src/lib.rs
@@ -133,14 +133,23 @@ mod arrow_registry;
 #[allow(clippy::unimplemented)]
 mod codegen;
 #[allow(clippy::unimplemented)]
+mod format;
+#[allow(clippy::unimplemented)]
 mod objects;
+
 pub mod report;
 
+/// In-memory generated files.
+///
+/// Generated by the codegen pass, modified in-place by thge post-processing passes (formatting
+/// etc), and finally written to disk by the I/O pass.
+pub type GeneratedFiles = std::collections::BTreeMap<camino::Utf8PathBuf, String>;
+
 pub use self::arrow_registry::{ArrowRegistry, LazyDatatype, LazyField};
 pub use self::codegen::{
-    CodeGenerator, CppCodeGenerator, DocsCodeGenerator, GeneratedFiles, PythonCodeGenerator,
-    RustCodeGenerator,
+    CodeGenerator, CppCodeGenerator, DocsCodeGenerator, PythonCodeGenerator, RustCodeGenerator,
 };
+pub use self::format::{CodeFormatter, CppCodeFormatter, PythonCodeFormatter, RustCodeFormatter};
 pub use self::objects::{
     Attributes, Docs, ElementType, Object, ObjectField, ObjectKind, ObjectSpecifics, Objects, Type,
 };

From 8e5411e2a5fb2e6e0a7152d59906c05265a9664a Mon Sep 17 00:00:00 2001
From: Clement Rey <cr.rey.clement@gmail.com>
Date: Wed, 18 Oct 2023 16:43:33 +0200
Subject: [PATCH 02/12] cpp formatter

---
 crates/re_types_builder/src/format/cpp.rs | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/crates/re_types_builder/src/format/cpp.rs b/crates/re_types_builder/src/format/cpp.rs
index e282f7734047..8007e54b17ee 100644
--- a/crates/re_types_builder/src/format/cpp.rs
+++ b/crates/re_types_builder/src/format/cpp.rs
@@ -5,7 +5,22 @@ use crate::CodeFormatter;
 pub struct CppCodeFormatter;
 
 impl CodeFormatter for CppCodeFormatter {
-    fn format(&mut self, reporter: &crate::Reporter, files: &mut crate::GeneratedFiles) {
-        todo!()
+    fn format(&mut self, _reporter: &crate::Reporter, files: &mut crate::GeneratedFiles) {
+        use rayon::prelude::*;
+
+        re_tracing::profile_function!();
+
+        files.par_iter_mut().for_each(|(filepath, contents)| {
+            *contents = if matches!(filepath.extension(), Some("cpp" | "hpp")) {
+                format_code(contents)
+            } else {
+                contents.clone()
+            };
+        });
     }
 }
+
+fn format_code(code: &str) -> String {
+    clang_format::clang_format_with_style(code, &clang_format::ClangFormatStyle::File)
+        .expect("Failed to run clang-format")
+}

From 1f1cef45e021cd021e120d043c7c18bf95ef505f Mon Sep 17 00:00:00 2001
From: Clement Rey <cr.rey.clement@gmail.com>
Date: Wed, 18 Oct 2023 16:48:30 +0200
Subject: [PATCH 03/12] rust formatter

---
 crates/re_types_builder/src/format/rust.rs | 41 ++++++++++++++++++++--
 1 file changed, 39 insertions(+), 2 deletions(-)

diff --git a/crates/re_types_builder/src/format/rust.rs b/crates/re_types_builder/src/format/rust.rs
index bf0c3b7b7bb7..dc642c60aeda 100644
--- a/crates/re_types_builder/src/format/rust.rs
+++ b/crates/re_types_builder/src/format/rust.rs
@@ -5,7 +5,44 @@ use crate::CodeFormatter;
 pub struct RustCodeFormatter;
 
 impl CodeFormatter for RustCodeFormatter {
-    fn format(&mut self, reporter: &crate::Reporter, files: &mut crate::GeneratedFiles) {
-        todo!()
+    fn format(&mut self, _reporter: &crate::Reporter, files: &mut crate::GeneratedFiles) {
+        use rayon::prelude::*;
+
+        re_tracing::profile_function!();
+
+        files.par_iter_mut().for_each(|(filepath, contents)| {
+            *contents = if matches!(filepath.extension(), Some("rs")) {
+                format_code(contents)
+            } else {
+                contents.clone()
+            };
+        });
     }
 }
+
+fn format_code(contents: &str) -> String {
+    re_tracing::profile_function!();
+
+    let mut contents = contents.replace(" :: ", "::"); // Fix `bytemuck :: Pod` -> `bytemuck::Pod`.
+
+    // Even though we already have used `prettyplease` we also
+    // need to run `cargo fmt`, since it catches some things `prettyplease` missed.
+    // We need to run `cago fmt` several times because it is not idempotent;
+    // see https://github.com/rust-lang/rustfmt/issues/5824
+    for _ in 0..2 {
+        // NOTE: We're purposefully ignoring the error here.
+        //
+        // In the very unlikely chance that the user doesn't have the `fmt` component installed,
+        // there's still no good reason to fail the build.
+        //
+        // The CI will catch the unformatted file at PR time and complain appropriately anyhow.
+
+        re_tracing::profile_scope!("rust-fmt");
+        use rust_format::Formatter as _;
+        if let Ok(formatted) = rust_format::RustFmt::default().format_str(&contents) {
+            contents = formatted;
+        }
+    }
+
+    contents
+}

From cf3825fae95f493c2527198825b46c506e9152dd Mon Sep 17 00:00:00 2001
From: Clement Rey <cr.rey.clement@gmail.com>
Date: Wed, 18 Oct 2023 16:54:00 +0200
Subject: [PATCH 04/12] python formatter

---
 crates/re_types_builder/src/format/python.rs | 140 ++++++++++++++++++-
 1 file changed, 137 insertions(+), 3 deletions(-)

diff --git a/crates/re_types_builder/src/format/python.rs b/crates/re_types_builder/src/format/python.rs
index c056f89ec4df..c77416ed4302 100644
--- a/crates/re_types_builder/src/format/python.rs
+++ b/crates/re_types_builder/src/format/python.rs
@@ -1,11 +1,145 @@
+use anyhow::Context;
+use camino::{Utf8Path, Utf8PathBuf};
+
 use crate::CodeFormatter;
 
 // ---
 
-pub struct PythonCodeFormatter;
+pub struct PythonCodeFormatter {
+    pub pkg_path: Utf8PathBuf,
+    pub testing_pkg_path: Utf8PathBuf,
+}
+
+impl PythonCodeFormatter {
+    pub fn new(pkg_path: impl Into<Utf8PathBuf>, testing_pkg_path: impl Into<Utf8PathBuf>) -> Self {
+        Self {
+            pkg_path: pkg_path.into(),
+            testing_pkg_path: testing_pkg_path.into(),
+        }
+    }
+}
 
 impl CodeFormatter for PythonCodeFormatter {
-    fn format(&mut self, reporter: &crate::Reporter, files: &mut crate::GeneratedFiles) {
-        todo!()
+    fn format(&mut self, _reporter: &crate::Reporter, files: &mut crate::GeneratedFiles) {
+        use rayon::prelude::*;
+
+        re_tracing::profile_function!();
+
+        // Running `black` once for each file is very slow, so we write all
+        // files to a temporary folder, format it, and copy back the results.
+
+        let tempdir = tempfile::tempdir().unwrap();
+        let tempdir_path = Utf8PathBuf::try_from(tempdir.path().to_owned()).unwrap();
+
+        files.par_iter().for_each(|(filepath, contents)| {
+            let formatted_source_path = format_path_for_tmp_dir(
+                &self.pkg_path,
+                &self.testing_pkg_path,
+                filepath,
+                &tempdir_path,
+            );
+            crate::codegen::common::write_file(&formatted_source_path, contents);
+        });
+
+        format_python_dir(&tempdir_path).unwrap();
+
+        // Read back and copy to the final destination:
+        files.par_iter_mut().for_each(|(filepath, contents)| {
+            let formatted_source_path = format_path_for_tmp_dir(
+                &self.pkg_path,
+                &self.testing_pkg_path,
+                filepath,
+                &tempdir_path,
+            );
+            *contents = std::fs::read_to_string(formatted_source_path).unwrap();
+        });
+    }
+}
+
+fn format_path_for_tmp_dir(
+    pkg_path: &Utf8Path,
+    testing_pkg_path: &Utf8Path,
+    filepath: &Utf8Path,
+    tempdir_path: &Utf8Path,
+) -> Utf8PathBuf {
+    // If the prefix is pkg_path, strip it, and then append to tempdir
+    // However, if the prefix is testing_pkg_path, strip it and insert an extra
+    // "testing" to avoid name collisions.
+    filepath.strip_prefix(pkg_path).map_or_else(
+        |_| {
+            tempdir_path
+                .join("testing")
+                .join(filepath.strip_prefix(testing_pkg_path).unwrap())
+        },
+        |f| tempdir_path.join(f),
+    )
+}
+
+fn format_python_dir(dir: &Utf8PathBuf) -> anyhow::Result<()> {
+    re_tracing::profile_function!();
+
+    // The order below is important and sadly we need to call black twice. Ruff does not yet
+    // fix line-length (See: https://github.com/astral-sh/ruff/issues/1904).
+    //
+    // 1) Call black, which among others things fixes line-length
+    // 2) Call ruff, which requires line-lengths to be correct
+    // 3) Call black again to cleanup some whitespace issues ruff might introduce
+
+    run_black_on_dir(dir).context("black")?;
+    run_ruff_on_dir(dir).context("ruff")?;
+    run_black_on_dir(dir).context("black")?;
+    Ok(())
+}
+
+fn python_project_path() -> Utf8PathBuf {
+    let path = crate::rerun_workspace_path()
+        .join("rerun_py")
+        .join("pyproject.toml");
+    assert!(path.exists(), "Failed to find {path:?}");
+    path
+}
+
+fn run_black_on_dir(dir: &Utf8PathBuf) -> anyhow::Result<()> {
+    re_tracing::profile_function!();
+    use std::process::{Command, Stdio};
+
+    let proc = Command::new("black")
+        .arg(format!("--config={}", python_project_path()))
+        .arg(dir)
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()?;
+
+    let output = proc.wait_with_output()?;
+
+    if output.status.success() {
+        Ok(())
+    } else {
+        let stdout = String::from_utf8_lossy(&output.stdout).trim().to_owned();
+        let stderr = String::from_utf8_lossy(&output.stderr).trim().to_owned();
+        anyhow::bail!("{stdout}\n{stderr}")
+    }
+}
+
+fn run_ruff_on_dir(dir: &Utf8PathBuf) -> anyhow::Result<()> {
+    re_tracing::profile_function!();
+    use std::process::{Command, Stdio};
+
+    let proc = Command::new("ruff")
+        .arg(format!("--config={}", python_project_path()))
+        .arg("--fix")
+        .arg(dir)
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()?;
+
+    let output = proc.wait_with_output()?;
+
+    if output.status.success() {
+        Ok(())
+    } else {
+        let stdout = String::from_utf8_lossy(&output.stdout).trim().to_owned();
+        let stderr = String::from_utf8_lossy(&output.stderr).trim().to_owned();
+        anyhow::bail!("{stdout}\n{stderr}")
     }
 }

From 3704b16f6d3ebb9446aa141ad2a8cf5f92a62615 Mon Sep 17 00:00:00 2001
From: Clement Rey <cr.rey.clement@gmail.com>
Date: Wed, 18 Oct 2023 17:18:59 +0200
Subject: [PATCH 05/12] add new generic generate_code and reimpl everything in
 terms of that

---
 crates/re_types_builder/src/format/mod.rs |   6 +
 crates/re_types_builder/src/lib.rs        | 303 ++++++----------------
 2 files changed, 91 insertions(+), 218 deletions(-)

diff --git a/crates/re_types_builder/src/format/mod.rs b/crates/re_types_builder/src/format/mod.rs
index ab274aeff520..7750ad81b594 100644
--- a/crates/re_types_builder/src/format/mod.rs
+++ b/crates/re_types_builder/src/format/mod.rs
@@ -4,6 +4,12 @@ pub trait CodeFormatter {
     fn format(&mut self, reporter: &crate::Reporter, files: &mut crate::GeneratedFiles);
 }
 
+pub struct NoopCodeFormatter;
+
+impl CodeFormatter for NoopCodeFormatter {
+    fn format(&mut self, _reporter: &crate::Reporter, _files: &mut crate::GeneratedFiles) {}
+}
+
 // ---
 
 mod cpp;
diff --git a/crates/re_types_builder/src/lib.rs b/crates/re_types_builder/src/lib.rs
index be6dfae8345b..ec6f5200e4bb 100644
--- a/crates/re_types_builder/src/lib.rs
+++ b/crates/re_types_builder/src/lib.rs
@@ -113,13 +113,15 @@
 )]
 mod reflection;
 
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 
 use anyhow::Context as _;
 use re_build_tools::{
     compute_crate_hash, compute_dir_filtered_hash, compute_dir_hash, compute_strings_hash,
 };
 
+use crate::format::NoopCodeFormatter;
+
 pub use self::reflection::reflection::{
     root_as_schema, BaseType as FbsBaseType, Enum as FbsEnum, EnumVal as FbsEnumVal,
     Field as FbsField, KeyValue as FbsKeyValue, Object as FbsObject, Schema as FbsSchema,
@@ -363,6 +365,44 @@ pub fn compute_re_types_hash(locations: &SourceLocations<'_>) -> String {
     new_hash
 }
 
+/// Generates, formats and optionally writes code.
+///
+/// Panics on error.
+fn generate_code(
+    reporter: &Reporter,
+    objects: &Objects,
+    arrow_registry: &ArrowRegistry,
+    generator: &mut dyn CodeGenerator,
+    formatter: &mut dyn CodeFormatter,
+    orphan_paths_opt_out: &BTreeSet<Utf8PathBuf>,
+) {
+    // 1. Generate in-memory code files.
+    let mut files = generator.generate(reporter, objects, arrow_registry);
+
+    // 2. Generate in-memory gitattribute files.
+    generate_gitattributes_for_generated_files(&mut files);
+
+    // 3. Format in-memory files.
+    formatter.format(reporter, &mut files);
+
+    // 4. Write all files to filesytem.
+    {
+        use rayon::prelude::*;
+
+        re_tracing::profile_scope!("write_files");
+
+        files.par_iter().for_each(|(filepath, contents)| {
+            crate::codegen::common::write_file(filepath, contents);
+        });
+    }
+
+    // 4. Remove orphaned files.
+    for path in orphan_paths_opt_out {
+        files.retain(|filepath, _| filepath.parent() != Some(path));
+    }
+    crate::codegen::common::remove_orphaned_files(reporter, &files);
+}
+
 /// Generates C++ code.
 ///
 /// Panics on error.
@@ -391,47 +431,28 @@ pub fn generate_cpp_code(
 ) {
     re_tracing::profile_function!();
 
-    // 1. Generate code files.
-    let mut gen = CppCodeGenerator::new(output_path.as_ref());
-    let mut files = gen.generate(reporter, objects, arrow_registry);
-    // 2. Generate attribute files.
-    generate_gitattributes_for_generated_files(&mut files);
-    // 3. Write all files.
-    {
-        use rayon::prelude::*;
+    let mut generator = CppCodeGenerator::new(output_path.as_ref());
+    let mut formatter = CppCodeFormatter;
 
-        re_tracing::profile_scope!("write_files");
-
-        files.par_iter().for_each(|(filepath, contents)| {
-            // There's more than cpp/hpp files in here, don't run clang-format on them!
-            let contents = if matches!(filepath.extension(), Some("cpp" | "hpp")) {
-                format_code(contents)
-            } else {
-                contents.clone()
-            };
-            crate::codegen::common::write_file(filepath, &contents);
-        });
-    }
-    // 4. Remove orphaned files.
     // NOTE: In rerun_cpp we have a directory where we share generated code with handwritten code.
     // Make sure to filter out that directory, or else we will end up removing those handwritten
     // files.
-    let root_src = output_path.as_ref().join("src/rerun");
-    files.retain(|filepath, _| filepath.parent() != Some(root_src.as_path()));
-    crate::codegen::common::remove_orphaned_files(reporter, &files);
-
-    fn format_code(code: &str) -> String {
-        clang_format::clang_format_with_style(code, &clang_format::ClangFormatStyle::File)
-            .expect("Failed to run clang-format")
-    }
+    let orphan_path_opt_out = output_path.as_ref().join("src/rerun");
+
+    generate_code(
+        reporter,
+        objects,
+        arrow_registry,
+        &mut generator,
+        &mut formatter,
+        &std::iter::once(orphan_path_opt_out).collect(),
+    );
 }
 
 /// Generates Rust code.
 ///
 /// Panics on error.
 ///
-/// - `output_crate_path`: path to the root of the output crate.
-///
 /// E.g.:
 /// ```no_run
 /// let (objects, arrow_registry) = re_types_builder::generate_lang_agnostic(
@@ -454,52 +475,17 @@ pub fn generate_rust_code(
 ) {
     re_tracing::profile_function!();
 
-    // 1. Generate code files.
-    let mut gen = RustCodeGenerator::new(workspace_path);
-    let mut files = gen.generate(reporter, objects, arrow_registry);
-    // 2. Generate attribute files.
-    generate_gitattributes_for_generated_files(&mut files);
-    // 3. Write all files.
-    write_files(&files);
-    // 4. Remove orphaned files.
-    crate::codegen::common::remove_orphaned_files(reporter, &files);
-
-    fn write_files(files: &GeneratedFiles) {
-        use rayon::prelude::*;
-
-        re_tracing::profile_function!();
-
-        files.par_iter().for_each(|(path, source)| {
-            write_file(path, source.clone());
-        });
-    }
+    let mut generator = RustCodeGenerator::new(workspace_path);
+    let mut formatter = RustCodeFormatter;
 
-    fn write_file(filepath: &Utf8PathBuf, mut contents: String) {
-        re_tracing::profile_function!();
-
-        contents = contents.replace(" :: ", "::"); // Fix `bytemuck :: Pod` -> `bytemuck::Pod`.
-
-        // Even though we already have used `prettyplease` we also
-        // need to run `cargo fmt`, since it catches some things `prettyplease` missed.
-        // We need to run `cago fmt` several times because it is not idempotent;
-        // see https://github.com/rust-lang/rustfmt/issues/5824
-        for _ in 0..2 {
-            // NOTE: We're purposefully ignoring the error here.
-            //
-            // In the very unlikely chance that the user doesn't have the `fmt` component installed,
-            // there's still no good reason to fail the build.
-            //
-            // The CI will catch the unformatted file at PR time and complain appropriately anyhow.
-
-            re_tracing::profile_scope!("rust-fmt");
-            use rust_format::Formatter as _;
-            if let Ok(formatted) = rust_format::RustFmt::default().format_str(&contents) {
-                contents = formatted;
-            }
-        }
-
-        crate::codegen::common::write_file(filepath, &contents);
-    }
+    generate_code(
+        reporter,
+        objects,
+        arrow_registry,
+        &mut generator,
+        &mut formatter,
+        &Default::default(),
+    );
 }
 
 /// Generates Python code.
@@ -532,132 +518,19 @@ pub fn generate_python_code(
 ) {
     re_tracing::profile_function!();
 
-    // 1. Generate code files.
-    let mut gen =
+    let mut generator =
         PythonCodeGenerator::new(output_pkg_path.as_ref(), testing_output_pkg_path.as_ref());
-    let mut files = gen.generate(reporter, objects, arrow_registry);
-    // 2. Generate attribute files.
-    generate_gitattributes_for_generated_files(&mut files);
-    // 3. Write all files.
-    write_files(&gen.pkg_path, &gen.testing_pkg_path, &files);
-    // 4. Remove orphaned files.
-    crate::codegen::common::remove_orphaned_files(reporter, &files);
-
-    fn write_files(pkg_path: &Utf8Path, testing_pkg_path: &Utf8Path, files: &GeneratedFiles) {
-        use rayon::prelude::*;
-
-        re_tracing::profile_function!();
-
-        // Running `black` once for each file is very slow, so we write all
-        // files to a temporary folder, format it, and copy back the results.
-
-        let tempdir = tempfile::tempdir().unwrap();
-        let tempdir_path = Utf8PathBuf::try_from(tempdir.path().to_owned()).unwrap();
-
-        files.par_iter().for_each(|(filepath, source)| {
-            let formatted_source_path =
-                format_path_for_tmp_dir(pkg_path, testing_pkg_path, filepath, &tempdir_path);
-            crate::codegen::common::write_file(&formatted_source_path, source);
-        });
-
-        format_python_dir(&tempdir_path).unwrap();
-
-        // Read back and copy to the final destination:
-        files.par_iter().for_each(|(filepath, _original_source)| {
-            let formatted_source_path =
-                format_path_for_tmp_dir(pkg_path, testing_pkg_path, filepath, &tempdir_path);
-            let formatted_source = std::fs::read_to_string(formatted_source_path).unwrap();
-            crate::codegen::common::write_file(filepath, &formatted_source);
-        });
-    }
-
-    fn format_path_for_tmp_dir(
-        pkg_path: &Utf8Path,
-        testing_pkg_path: &Utf8Path,
-        filepath: &Utf8Path,
-        tempdir_path: &Utf8Path,
-    ) -> Utf8PathBuf {
-        // If the prefix is pkg_path, strip it, and then append to tempdir
-        // However, if the prefix is testing_pkg_path, strip it and insert an extra
-        // "testing" to avoid name collisions.
-        filepath.strip_prefix(pkg_path).map_or_else(
-            |_| {
-                tempdir_path
-                    .join("testing")
-                    .join(filepath.strip_prefix(testing_pkg_path).unwrap())
-            },
-            |f| tempdir_path.join(f),
-        )
-    }
-
-    fn format_python_dir(dir: &Utf8PathBuf) -> anyhow::Result<()> {
-        re_tracing::profile_function!();
-
-        // The order below is important and sadly we need to call black twice. Ruff does not yet
-        // fix line-length (See: https://github.com/astral-sh/ruff/issues/1904).
-        //
-        // 1) Call black, which among others things fixes line-length
-        // 2) Call ruff, which requires line-lengths to be correct
-        // 3) Call black again to cleanup some whitespace issues ruff might introduce
-
-        run_black_on_dir(dir).context("black")?;
-        run_ruff_on_dir(dir).context("ruff")?;
-        run_black_on_dir(dir).context("black")?;
-        Ok(())
-    }
-
-    fn python_project_path() -> Utf8PathBuf {
-        let path = crate::rerun_workspace_path()
-            .join("rerun_py")
-            .join("pyproject.toml");
-        assert!(path.exists(), "Failed to find {path:?}");
-        path
-    }
-
-    fn run_black_on_dir(dir: &Utf8PathBuf) -> anyhow::Result<()> {
-        re_tracing::profile_function!();
-        use std::process::{Command, Stdio};
-
-        let proc = Command::new("black")
-            .arg(format!("--config={}", python_project_path()))
-            .arg(dir)
-            .stdout(Stdio::piped())
-            .stderr(Stdio::piped())
-            .spawn()?;
-
-        let output = proc.wait_with_output()?;
-
-        if output.status.success() {
-            Ok(())
-        } else {
-            let stdout = String::from_utf8_lossy(&output.stdout).trim().to_owned();
-            let stderr = String::from_utf8_lossy(&output.stderr).trim().to_owned();
-            anyhow::bail!("{stdout}\n{stderr}")
-        }
-    }
-
-    fn run_ruff_on_dir(dir: &Utf8PathBuf) -> anyhow::Result<()> {
-        re_tracing::profile_function!();
-        use std::process::{Command, Stdio};
-
-        let proc = Command::new("ruff")
-            .arg(format!("--config={}", python_project_path()))
-            .arg("--fix")
-            .arg(dir)
-            .stdout(Stdio::piped())
-            .stderr(Stdio::piped())
-            .spawn()?;
-
-        let output = proc.wait_with_output()?;
-
-        if output.status.success() {
-            Ok(())
-        } else {
-            let stdout = String::from_utf8_lossy(&output.stdout).trim().to_owned();
-            let stderr = String::from_utf8_lossy(&output.stderr).trim().to_owned();
-            anyhow::bail!("{stdout}\n{stderr}")
-        }
-    }
+    let mut formatter =
+        PythonCodeFormatter::new(output_pkg_path.as_ref(), testing_output_pkg_path.as_ref());
+
+    generate_code(
+        reporter,
+        objects,
+        arrow_registry,
+        &mut generator,
+        &mut formatter,
+        &Default::default(),
+    );
 }
 
 pub fn generate_docs(
@@ -670,23 +543,17 @@ pub fn generate_docs(
 
     re_log::info!("Generating docs to {}", output_docs_dir.as_ref());
 
-    // 1. Generate code files.
-    let mut gen = DocsCodeGenerator::new(output_docs_dir.as_ref());
-    let mut files = gen.generate(reporter, objects, arrow_registry);
-    // 2. Generate attribute files.
-    generate_gitattributes_for_generated_files(&mut files);
-    // 3. Write all files.
-    {
-        use rayon::prelude::*;
-
-        re_tracing::profile_scope!("write_files");
+    let mut generator = DocsCodeGenerator::new(output_docs_dir.as_ref());
+    let mut formatter = NoopCodeFormatter;
 
-        files.par_iter().for_each(|(filepath, contents)| {
-            crate::codegen::common::write_file(filepath, contents);
-        });
-    }
-    // 4. Remove orphaned files.
-    crate::codegen::common::remove_orphaned_files(reporter, &files);
+    generate_code(
+        reporter,
+        objects,
+        arrow_registry,
+        &mut generator,
+        &mut formatter,
+        &Default::default(),
+    );
 }
 
 pub(crate) fn rerun_workspace_path() -> camino::Utf8PathBuf {

From aa3ac7e8d488386bb8e74a08c5ee0f481f1db994 Mon Sep 17 00:00:00 2001
From: Clement Rey <cr.rey.clement@gmail.com>
Date: Wed, 18 Oct 2023 17:27:52 +0200
Subject: [PATCH 06/12] implement --check support

---
 .../src/bin/build_re_types.rs                 | 16 +++++++-----
 crates/re_types_builder/src/lib.rs            | 25 +++++++++++++++++--
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/crates/re_types_builder/src/bin/build_re_types.rs b/crates/re_types_builder/src/bin/build_re_types.rs
index 7c4ba263f4c5..c298ae2cef63 100644
--- a/crates/re_types_builder/src/bin/build_re_types.rs
+++ b/crates/re_types_builder/src/bin/build_re_types.rs
@@ -44,6 +44,7 @@ fn main() {
 
     let mut profiler = re_tracing::Profiler::default();
     let mut always_run = false;
+    let mut check = false;
 
     for arg in std::env::args().skip(1) {
         match arg.as_str() {
@@ -51,10 +52,9 @@ fn main() {
                 println!("Usage: [--help] [--force] [--profile]");
                 return;
             }
-            "--force" => {
-                always_run = true;
-            }
+            "--force" => always_run = true,
             "--profile" => profiler.start(),
+            "--check" => check = true,
             _ => {
                 eprintln!("Unknown argument: {arg:?}");
                 return;
@@ -114,13 +114,15 @@ fn main() {
             &reporter,
             cpp_output_dir_path,
             &objects,
-            &arrow_registry
+            &arrow_registry,
+            check,
         ),
         || re_types_builder::generate_rust_code(
             &reporter,
             workspace_dir,
             &objects,
-            &arrow_registry
+            &arrow_registry,
+            check,
         ),
         || re_types_builder::generate_python_code(
             &reporter,
@@ -128,12 +130,14 @@ fn main() {
             python_testing_output_dir_path,
             &objects,
             &arrow_registry,
+            check,
         ),
         || re_types_builder::generate_docs(
             &reporter,
             docs_content_dir_path,
             &objects,
-            &arrow_registry
+            &arrow_registry,
+            check,
         ),
     );
 
diff --git a/crates/re_types_builder/src/lib.rs b/crates/re_types_builder/src/lib.rs
index ec6f5200e4bb..2349fd970d82 100644
--- a/crates/re_types_builder/src/lib.rs
+++ b/crates/re_types_builder/src/lib.rs
@@ -367,6 +367,8 @@ pub fn compute_re_types_hash(locations: &SourceLocations<'_>) -> String {
 
 /// Generates, formats and optionally writes code.
 ///
+/// If `check` is true, this will run a comparison check instead of writing files to disk.
+///
 /// Panics on error.
 fn generate_code(
     reporter: &Reporter,
@@ -375,7 +377,10 @@ fn generate_code(
     generator: &mut dyn CodeGenerator,
     formatter: &mut dyn CodeFormatter,
     orphan_paths_opt_out: &BTreeSet<Utf8PathBuf>,
+    check: bool,
 ) {
+    use rayon::prelude::*;
+
     // 1. Generate in-memory code files.
     let mut files = generator.generate(reporter, objects, arrow_registry);
 
@@ -385,10 +390,18 @@ fn generate_code(
     // 3. Format in-memory files.
     formatter.format(reporter, &mut files);
 
+    if check {
+        files.par_iter().for_each(|(filepath, contents)| {
+            let cur_contents = std::fs::read_to_string(filepath).unwrap();
+            if *contents != cur_contents {
+                reporter.error(filepath.as_str(), "", "out of sync");
+            }
+        });
+        return;
+    }
+
     // 4. Write all files to filesytem.
     {
-        use rayon::prelude::*;
-
         re_tracing::profile_scope!("write_files");
 
         files.par_iter().for_each(|(filepath, contents)| {
@@ -428,6 +441,7 @@ pub fn generate_cpp_code(
     output_path: impl AsRef<Utf8Path>,
     objects: &Objects,
     arrow_registry: &ArrowRegistry,
+    check: bool,
 ) {
     re_tracing::profile_function!();
 
@@ -446,6 +460,7 @@ pub fn generate_cpp_code(
         &mut generator,
         &mut formatter,
         &std::iter::once(orphan_path_opt_out).collect(),
+        check,
     );
 }
 
@@ -472,6 +487,7 @@ pub fn generate_rust_code(
     workspace_path: impl Into<Utf8PathBuf>,
     objects: &Objects,
     arrow_registry: &ArrowRegistry,
+    check: bool,
 ) {
     re_tracing::profile_function!();
 
@@ -485,6 +501,7 @@ pub fn generate_rust_code(
         &mut generator,
         &mut formatter,
         &Default::default(),
+        check,
     );
 }
 
@@ -515,6 +532,7 @@ pub fn generate_python_code(
     testing_output_pkg_path: impl AsRef<Utf8Path>,
     objects: &Objects,
     arrow_registry: &ArrowRegistry,
+    check: bool,
 ) {
     re_tracing::profile_function!();
 
@@ -530,6 +548,7 @@ pub fn generate_python_code(
         &mut generator,
         &mut formatter,
         &Default::default(),
+        check,
     );
 }
 
@@ -538,6 +557,7 @@ pub fn generate_docs(
     output_docs_dir: impl AsRef<Utf8Path>,
     objects: &Objects,
     arrow_registry: &ArrowRegistry,
+    check: bool,
 ) {
     re_tracing::profile_function!();
 
@@ -553,6 +573,7 @@ pub fn generate_docs(
         &mut generator,
         &mut formatter,
         &Default::default(),
+        check,
     );
 }
 

From 171d36ab48b245e335b394af0aeb554c4a72f561 Mon Sep 17 00:00:00 2001
From: Clement Rey <cr.rey.clement@gmail.com>
Date: Wed, 18 Oct 2023 17:30:03 +0200
Subject: [PATCH 07/12] use --check on CI

---
 .github/workflows/contrib_checks.yml  | 6 +-----
 .github/workflows/reusable_checks.yml | 6 +-----
 2 files changed, 2 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/contrib_checks.yml b/.github/workflows/contrib_checks.yml
index d6a690c75567..db87027c139c 100644
--- a/.github/workflows/contrib_checks.yml
+++ b/.github/workflows/contrib_checks.yml
@@ -92,11 +92,7 @@ jobs:
         uses: actions-rs/cargo@v1
         with:
           command: codegen
-          args: --force
-
-      - name: No Diffs From Running Codegen
-        run: |
-          git diff --exit-code
+          args: --force --check
 
   rs-lints:
     name: Rust lints (fmt, check, cranky, tests, doc)
diff --git a/.github/workflows/reusable_checks.yml b/.github/workflows/reusable_checks.yml
index cd83875c7e3b..c5aca13f41b0 100644
--- a/.github/workflows/reusable_checks.yml
+++ b/.github/workflows/reusable_checks.yml
@@ -134,11 +134,7 @@ jobs:
         uses: actions-rs/cargo@v1
         with:
           command: codegen
-          args: --force
-
-      - name: No Diffs From Running Codegen
-        run: |
-          git diff --exit-code
+          args: --force --check
 
   rs-lints:
     name: Rust lints (fmt, check, cranky, tests, doc)

From 7b17fed3b312e2c60d04228baa9c0b823b75fc1b Mon Sep 17 00:00:00 2001
From: Clement Rey <cr.rey.clement@gmail.com>
Date: Wed, 18 Oct 2023 17:30:22 +0200
Subject: [PATCH 08/12] show that it works!

---
 crates/re_types/definitions/rerun/datatypes/uvec2d.fbs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/re_types/definitions/rerun/datatypes/uvec2d.fbs b/crates/re_types/definitions/rerun/datatypes/uvec2d.fbs
index e16fcbd40776..185832ee9f82 100644
--- a/crates/re_types/definitions/rerun/datatypes/uvec2d.fbs
+++ b/crates/re_types/definitions/rerun/datatypes/uvec2d.fbs
@@ -7,7 +7,7 @@ namespace rerun.datatypes;
 
 // ---
 
-/// A uint32 vector in 2D space.
+/// A uint32 vector in 2D space LOL.
 struct UVec2D (
   "attr.arrow.transparent",
   "attr.python.aliases": "npt.NDArray[Any], npt.ArrayLike, Sequence[int]",

From 0f58b572b47d018287f72b5fcc2a27f889e2799a Mon Sep 17 00:00:00 2001
From: Clement Rey <cr.rey.clement@gmail.com>
Date: Wed, 18 Oct 2023 17:33:09 +0200
Subject: [PATCH 09/12] lints

---
 crates/re_types_builder/src/lib.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/crates/re_types_builder/src/lib.rs b/crates/re_types_builder/src/lib.rs
index 2349fd970d82..1a3bf0c6b366 100644
--- a/crates/re_types_builder/src/lib.rs
+++ b/crates/re_types_builder/src/lib.rs
@@ -143,7 +143,7 @@ pub mod report;
 
 /// In-memory generated files.
 ///
-/// Generated by the codegen pass, modified in-place by thge post-processing passes (formatting
+/// Generated by the codegen pass, modified in-place by the post-processing passes (formatting
 /// etc), and finally written to disk by the I/O pass.
 pub type GeneratedFiles = std::collections::BTreeMap<camino::Utf8PathBuf, String>;
 
@@ -400,7 +400,7 @@ fn generate_code(
         return;
     }
 
-    // 4. Write all files to filesytem.
+    // 4. Write all files to filesystem.
     {
         re_tracing::profile_scope!("write_files");
 

From f38107a4ae90d56855b1e9431a4286b701da3701 Mon Sep 17 00:00:00 2001
From: Clement Rey <cr.rey.clement@gmail.com>
Date: Wed, 18 Oct 2023 17:38:47 +0200
Subject: [PATCH 10/12] Revert "show that it works!"

This reverts commit dfbdfc9bd0602c757f05355ebcac6b0e3ea56611.
---
 crates/re_types/definitions/rerun/datatypes/uvec2d.fbs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/re_types/definitions/rerun/datatypes/uvec2d.fbs b/crates/re_types/definitions/rerun/datatypes/uvec2d.fbs
index 185832ee9f82..e16fcbd40776 100644
--- a/crates/re_types/definitions/rerun/datatypes/uvec2d.fbs
+++ b/crates/re_types/definitions/rerun/datatypes/uvec2d.fbs
@@ -7,7 +7,7 @@ namespace rerun.datatypes;
 
 // ---
 
-/// A uint32 vector in 2D space LOL.
+/// A uint32 vector in 2D space.
 struct UVec2D (
   "attr.arrow.transparent",
   "attr.python.aliases": "npt.NDArray[Any], npt.ArrayLike, Sequence[int]",

From e57741a55a36601cf83269291aba5f61b3fa8fc1 Mon Sep 17 00:00:00 2001
From: Clement Rey <cr.rey.clement@gmail.com>
Date: Wed, 18 Oct 2023 17:52:46 +0200
Subject: [PATCH 11/12] lints

---
 crates/re_types_builder/src/lib.rs | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/crates/re_types_builder/src/lib.rs b/crates/re_types_builder/src/lib.rs
index 1a3bf0c6b366..7a16abdfd770 100644
--- a/crates/re_types_builder/src/lib.rs
+++ b/crates/re_types_builder/src/lib.rs
@@ -418,6 +418,8 @@ fn generate_code(
 
 /// Generates C++ code.
 ///
+/// If `check` is true, this will run a comparison check instead of writing files to disk.
+///
 /// Panics on error.
 ///
 /// - `output_path`: path to the root of the output.
@@ -429,11 +431,13 @@ fn generate_code(
 ///     "./definitions/rerun/archetypes.fbs",
 /// );
 /// # let reporter = re_types_builder::report::init().1;
+/// # let check = false;
 /// re_types_builder::generate_cpp_code(
 ///     &reporter,
 ///     ".",
 ///     &objects,
 ///     &arrow_registry,
+///     check,
 /// );
 /// ```
 pub fn generate_cpp_code(
@@ -466,6 +470,8 @@ pub fn generate_cpp_code(
 
 /// Generates Rust code.
 ///
+/// If `check` is true, this will run a comparison check instead of writing files to disk.
+///
 /// Panics on error.
 ///
 /// E.g.:
@@ -475,11 +481,13 @@ pub fn generate_cpp_code(
 ///     "./definitions/rerun/archetypes.fbs",
 /// );
 /// # let reporter = re_types_builder::report::init().1;
+/// # let check = false;
 /// re_types_builder::generate_rust_code(
 ///     &reporter,
 ///     ".",
 ///     &objects,
 ///     &arrow_registry,
+///     check,
 /// );
 /// ```
 pub fn generate_rust_code(
@@ -507,6 +515,8 @@ pub fn generate_rust_code(
 
 /// Generates Python code.
 ///
+/// If `check` is true, this will run a comparison check instead of writing files to disk.
+///
 /// Panics on error.
 ///
 /// - `output_pkg_path`: path to the root of the output package.
@@ -518,12 +528,14 @@ pub fn generate_rust_code(
 ///     "./definitions/rerun/archetypes.fbs",
 /// );
 /// # let reporter = re_types_builder::report::init().1;
+/// # let check = false;
 /// re_types_builder::generate_python_code(
 ///     &reporter,
 ///     "./rerun_py/rerun_sdk",
 ///     "./rerun_py/tests",
 ///     &objects,
 ///     &arrow_registry,
+///     check,
 /// );
 /// ```
 pub fn generate_python_code(

From b02b166619ab1ede4cdcac149a42c392ca8f13a4 Mon Sep 17 00:00:00 2001
From: Clement Rey <cr.rey.clement@gmail.com>
Date: Thu, 19 Oct 2023 10:12:36 +0200
Subject: [PATCH 12/12] --check implies --force

---
 crates/re_types_builder/src/bin/build_re_types.rs | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/crates/re_types_builder/src/bin/build_re_types.rs b/crates/re_types_builder/src/bin/build_re_types.rs
index c298ae2cef63..84ae5263a59c 100644
--- a/crates/re_types_builder/src/bin/build_re_types.rs
+++ b/crates/re_types_builder/src/bin/build_re_types.rs
@@ -54,7 +54,10 @@ fn main() {
             }
             "--force" => always_run = true,
             "--profile" => profiler.start(),
-            "--check" => check = true,
+            "--check" => {
+                always_run = true;
+                check = true;
+            }
             _ => {
                 eprintln!("Unknown argument: {arg:?}");
                 return;