feat(mro): Improve martian_make_mro

martian/src/lib.rs

@@ -3,25 +3,23 @@
 //!
 //! ## Documentation
 //! For a guide style documentation and examples, visit: [https://martian-lang.github.io/martian-rust/](https://martian-lang.github.io/martian-rust/#/)
-//!
 
 pub use anyhow::Error;
-use anyhow::{format_err, Context};
+use anyhow::{ensure, Context, Result};
 use backtrace::Backtrace;
 use log::{error, info};
-
 use std::collections::HashMap;
 use std::fmt::Write as FmtWrite;
 use std::fs::File;
 use std::io::Write as IoWrite;
 use std::os::unix::io::{FromRawFd, IntoRawFd};
 use std::path::Path;
-
 use std::{io, panic};
 use time::format_description::modifier::{Day, Hour, Minute, Month, Second, Year};
 use time::format_description::FormatItem::Literal;
 use time::format_description::{Component, FormatItem};
 use time::OffsetDateTime;
+use utils::current_executable;
 
 mod metadata;
 pub use metadata::*;
@@ -42,15 +40,15 @@ pub use log::LevelFilter;
 pub use mro::*;
 pub mod prelude;
 
-pub fn initialize(args: Vec<String>) -> Result<Metadata, Error> {
+pub fn initialize(args: Vec<String>) -> Result<Metadata> {
     let mut md = Metadata::new(args);
     md.update_jobinfo()?;
 
     Ok(md)
 }
 
 #[cold]
-fn write_errors(msg: &str, is_assert: bool) -> Result<(), Error> {
+fn write_errors(msg: &str, is_assert: bool) -> Result<()> {
     let mut err_file: File = unsafe { File::from_raw_fd(4) };
 
     // We want to aggressively avoid allocations here if we can, since one
@@ -79,7 +77,7 @@ fn write_errors(msg: &str, is_assert: bool) -> Result<(), Error> {
 // We could use the proc macro, but then we'd need
 // to compile the proc macro crate, which would slow down build times
 // significantly for very little benefit in readability.
-pub(crate) const DATE_FORMAT: &[FormatItem] = &[
+pub(crate) const DATE_FORMAT: &[FormatItem<'_>] = &[
     FormatItem::Component(Component::Year(Year::default())),
     Literal(b"-"),
     FormatItem::Component(Component::Month(Month::default())),
@@ -195,13 +193,13 @@ fn martian_entry_point<S: std::hash::BuildHasher>(
     };
 
     // Get the stage implementation
-    let _stage = stage_map.get(&md.stage_name).ok_or_else(
+    let stage = stage_map.get(&md.stage_name).with_context(
         #[cold]
-        || format_err!("Couldn't find requested Martian stage: {}", md.stage_name),
+        || format!("Couldn't find requested Martian stage: {}", md.stage_name),
     );
 
     // special handler for non-existent stage
-    let stage = match _stage {
+    let stage = match stage {
         Ok(s) => s,
         Err(e) => {
             let _ = write_errors(&format!("{e:?}"), false);
@@ -286,48 +284,43 @@ fn report_error(md: &mut Metadata, e: &Error, is_assert: bool) {
     let _ = write_errors(&format!("{e:#}"), is_assert);
 }
 
+/// Return the environment variable CARGO_PKG_NAME or the current executable name.
 fn get_generator_name() -> String {
-    std::env::var("CARGO_BIN_NAME")
-        .or_else(|_| std::env::var("CARGO_CRATE_NAME"))
-        .or_else(|_| std::env::var("CARGO_PKG_NAME"))
-        .unwrap_or_else(|_| {
-            option_env!("CARGO_BIN_NAME")
-                .or(option_env!("CARGO_CRATE_NAME"))
-                .unwrap_or("martian-rust")
-                .into()
-        })
+    std::env::var("CARGO_PKG_NAME").unwrap_or_else(|_| current_executable())
 }
 
+/// Write MRO to filename or stdout.
 pub fn martian_make_mro(
     header_comment: &str,
-    file_name: Option<impl AsRef<Path>>,
+    filename: Option<impl AsRef<Path>>,
     rewrite: bool,
     mro_registry: Vec<StageMro>,
-) -> Result<(), Error> {
-    if let Some(ref f) = file_name {
-        let file_path = f.as_ref();
-        if file_path.is_dir() {
-            return Err(format_err!(
-                "Error! Path {} is a directory!",
-                file_path.display()
-            ));
-        }
-        if file_path.exists() && !rewrite {
-            return Err(format_err!(
-                "File {} exists. You need to explicitly mention if it is okay to rewrite.",
-                file_path.display()
-            ));
-        }
+) -> Result<()> {
+    if let Some(filename) = &filename {
+        let filename = filename.as_ref();
+        ensure!(
+            !filename.is_dir(),
+            "Path {} is a directory",
+            filename.display()
+        );
+        ensure!(
+            rewrite || !filename.exists(),
+            "File {} exists. Use --rewrite to overwrite it.",
+            filename.display()
+        );
     }
 
-    let final_mro_string = make_mro_string(header_comment, &mro_registry);
-    match file_name {
-        Some(f) => {
-            let mut output = File::create(f)?;
-            output.write_all(final_mro_string.as_bytes())?;
+    let mro = make_mro_string(header_comment, &mro_registry);
+    match filename {
+        Some(filename) => {
+            let filename = filename.as_ref();
+            File::create(filename)
+                .with_context(|| filename.display().to_string())?
+                .write_all(mro.as_bytes())
+                .with_context(|| filename.display().to_string())?;
         }
         None => {
-            print!("{final_mro_string}");
+            print!("{mro}");
         }
     }
     Ok(())
@@ -361,8 +354,7 @@ pub fn make_mro_string(header_comment: &str, mro_registry: &[StageMro]) -> Strin
             header_comment
                 .lines()
                 .all(|line| line.trim_end().is_empty() || line.starts_with('#')),
-            "All non-empty header lines must start with '#', but got\n{}",
-            header_comment
+            "All non-empty header lines must start with '#', but got\n{header_comment}"
         );
         format!(
             "{}

martian/src/metadata.rs

@@ -179,7 +179,7 @@ impl Metadata {
 
     /// Update the Martian journal -- so that Martian knows what we've updated
     fn update_journal(&self, name: &str) -> Result<()> {
-        let journal_name: Cow<str> = if self.stage_type != "main" {
+        let journal_name: Cow<'_, str> = if self.stage_type != "main" {
             format!("{}_{name}", self.stage_type).into()
         } else {
             name.into()

martian/src/utils.rs

@@ -52,12 +52,14 @@ pub fn to_camel_case(stage_name: &str) -> String {
 /// Parse the `env::args()` and return the name of the
 /// current executable as a String
 pub fn current_executable() -> String {
-    let args: Vec<_> = std::env::args().collect();
-    std::path::Path::new(&args[0])
+    Path::new(&std::env::args().next().unwrap())
+        .canonicalize()
+        .unwrap()
         .file_name()
         .unwrap()
-        .to_string_lossy()
-        .into_owned()
+        .to_str()
+        .unwrap()
+        .to_string()
 }
 
 /// Given a filename and an extension, return the filename with the correct extension.