From ec73c3f2629e910b6449012855d3e7d0f4c0f543 Mon Sep 17 00:00:00 2001
From: David Freese <freese@verily.com>
Date: Thu, 25 Feb 2021 23:27:43 -0800
Subject: [PATCH] Allow for the specification of the rustfmt path

In some situations, rustfmt can't be counted on to be in your PATH.  In
those cases, this adds a format_with(...) function to the builder to
specify the path to rustfmt, which is then passed into the new
format_with function.

A small dance is done here with fmt(...) and format_with(...) to
preserve the public API of the crate, so this, if I understand
correctly, should be an API-compatible change.

I'm not 100% this would be the way I would want to go, but I wanted to
put it up for discussion.  There are more roundabout approaches I could
take to format each file in our build system, but I had deferred from
doing that previously.   The relevant section of a public draft of the
implementation I had put together is here:

https://github.com/bazelbuild/rules_rust/pull/479/files#r583439508

Also open to suggestions as to where this should be tested.  I didn't
see, but may have missed, a good place for it.
---
 tonic-build/src/lib.rs   | 10 +++++++++-
 tonic-build/src/prost.rs | 20 +++++++++++++++++++-
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/tonic-build/src/lib.rs b/tonic-build/src/lib.rs
index 3d312a5e7..3adb97a1d 100644
--- a/tonic-build/src/lib.rs
+++ b/tonic-build/src/lib.rs
@@ -158,6 +158,14 @@ pub trait Method {
 #[cfg(feature = "rustfmt")]
 #[cfg_attr(docsrs, doc(cfg(feature = "rustfmt")))]
 pub fn fmt(out_dir: &str) {
+    // Since fmt is public, preserve the API of fmt by dispatching to a helper function that is
+    // called within the crate.
+    format_with("rustfmt", &out_dir);
+}
+
+#[cfg(feature = "rustfmt")]
+#[cfg_attr(docsrs, doc(cfg(feature = "rustfmt")))]
+fn format_with(rustfmt_path: impl AsRef<std::ffi::OsStr>, out_dir: &str) {
     let dir = std::fs::read_dir(out_dir).unwrap();
 
     for entry in dir {
@@ -165,7 +173,7 @@ pub fn fmt(out_dir: &str) {
         if !file.ends_with(".rs") {
             continue;
         }
-        let result = Command::new("rustfmt")
+        let result = Command::new(rustfmt_path.as_ref())
             .arg("--emit")
             .arg("files")
             .arg("--edition")
diff --git a/tonic-build/src/prost.rs b/tonic-build/src/prost.rs
index 49c4bfde3..14a80b7b8 100644
--- a/tonic-build/src/prost.rs
+++ b/tonic-build/src/prost.rs
@@ -21,6 +21,7 @@ pub fn configure() -> Builder {
         compile_well_known_types: false,
         #[cfg(feature = "rustfmt")]
         format: true,
+        rustfmt_path: None,
         emit_package: true,
     }
 }
@@ -210,6 +211,8 @@ pub struct Builder {
     out_dir: Option<PathBuf>,
     #[cfg(feature = "rustfmt")]
     format: bool,
+    #[cfg(feature = "rustfmt")]
+    rustfmt_path: Option<PathBuf>,
 }
 
 impl Builder {
@@ -239,6 +242,13 @@ impl Builder {
         self
     }
 
+    /// The path that should be used to find rustfmt, if format is enabled.
+    #[cfg(feature = "rustfmt")]
+    pub fn format_with(mut self, path: impl AsRef<Path>) -> Self {
+        self.rustfmt_path = Some(path.as_ref().to_path_buf());
+        self
+    }
+
     /// Set the output directory to generate code to.
     ///
     /// Defaults to the `OUT_DIR` environment variable.
@@ -331,6 +341,11 @@ impl Builder {
 
         #[cfg(feature = "rustfmt")]
         let format = self.format;
+        #[cfg(feature = "rustfmt")]
+        let rustfmt_path = self
+            .rustfmt_path
+            .clone()
+            .unwrap_or_else(|| "rustfmt".into());
 
         config.out_dir(out_dir.clone());
         if let Some(path) = self.file_descriptor_set_path.as_ref() {
@@ -355,7 +370,10 @@ impl Builder {
         #[cfg(feature = "rustfmt")]
         {
             if format {
-                super::fmt(out_dir.to_str().expect("Expected utf8 out_dir"));
+                super::format_with(
+                    rustfmt_path,
+                    out_dir.to_str().expect("Expected utf8 out_dir"),
+                );
             }
         }