Codegen: Multi-pass refactoring & support for --check (#3918)

- Adds a new `CodeFormatter` trait.
- Adds a new trait-based generic `generate_code` helper, and reimplement
all language-specific backends in terms of it.
- Adds `--check` flag.
- Use `--check` instead of git dirtyness checks on CI.

---

- Fixes #3495

.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)

.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)

crates/re_types_builder/src/bin/build_re_types.rs

@@ -44,17 +44,20 @@ 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() {
  "--help" => {
  println!("Usage: [--help] [--force] [--profile]");
  return;
  }
- "--force" => {
+ "--force" => always_run = true,
+ "--profile" => profiler.start(),
+ "--check" => {
  always_run = true;
+ check = true;
  }
- "--profile" => profiler.start(),
  _ => {
  eprintln!("Unknown argument: {arg:?}");
  return;
@@ -114,26 +117,30 @@ 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,
  python_output_dir_path,
  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,
  ),
  );
 

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;
 }
 
 // ---

crates/re_types_builder/src/format/cpp.rs

@@ -0,0 +1,26 @@
+use crate::CodeFormatter;
+
+// ---
+
+pub struct CppCodeFormatter;
+
+impl CodeFormatter for CppCodeFormatter {
+ 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")
+}

crates/re_types_builder/src/format/mod.rs

@@ -0,0 +1,21 @@
+/// Implements the formatting pass.
+pub trait CodeFormatter {
+ /// Formats generated files in-place.
+ 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;
+mod python;
+mod rust;
+
+pub use self::cpp::CppCodeFormatter;
+pub use self::python::PythonCodeFormatter;
+pub use self::rust::RustCodeFormatter;

crates/re_types_builder/src/format/python.rs

@@ -0,0 +1,145 @@
+use anyhow::Context;
+use camino::{Utf8Path, Utf8PathBuf};
+
+use crate::CodeFormatter;
+
+// ---
+
+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) {
+ 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}")
+ }
+}

crates/re_types_builder/src/format/rust.rs

@@ -0,0 +1,48 @@
+use crate::CodeFormatter;
+
+// ---
+
+pub struct RustCodeFormatter;
+
+impl CodeFormatter for RustCodeFormatter {
+ 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
+}