src/cargo-fmt/main.rs

   1 // Inspired by Paul Woolcock's cargo-fmt (https://github.com/pwoolcoc/cargo-fmt/).
   2
   3 #![cfg(not(test))]
   4 #![deny(warnings)]
   5
   6 use cargo_metadata;
   7 use getopts;
   8
   9 use std::cmp::Ordering;
  10 use std::collections::{BTreeMap, BTreeSet};
  11 use std::env;
  12 use std::fs;
  13 use std::hash::{Hash, Hasher};
  14 use std::io::{self, Write};
  15 use std::iter::FromIterator;
  16 use std::path::{Path, PathBuf};
  17 use std::process::Command;
  18 use std::str;
  19
  20 use getopts::{Matches, Options};
  21
  22 fn main() {
  23     let exit_status = execute();
  24     std::io::stdout().flush().unwrap();
  25     std::process::exit(exit_status);
  26 }
  27
  28 const SUCCESS: i32 = 0;
  29 const FAILURE: i32 = 1;
  30
  31 fn execute() -> i32 {
  32     let mut opts = getopts::Options::new();
  33     opts.optflag("h", "help", "show this message");
  34     opts.optflag("q", "quiet", "no output printed to stdout");
  35     opts.optflag("v", "verbose", "use verbose output");
  36     opts.optmulti(
  37         "p",
  38         "package",
  39         "specify package to format (only usable in workspaces)",
  40         "<package>",
  41     );
  42     opts.optflag("", "version", "print rustfmt version and exit");
  43     opts.optflag("", "all", "format all packages (only usable in workspaces)");
  44
  45     // If there is any invalid argument passed to `cargo fmt`, return without formatting.
  46     let mut is_package_arg = false;
  47     for arg in env::args().skip(2).take_while(|a| a != "--") {
  48         if arg.starts_with('-') {
  49             is_package_arg = arg.starts_with("--package") | arg.starts_with("-p");
  50         } else if !is_package_arg {
  51             print_usage_to_stderr(&opts, &format!("Invalid argument: `{}`.", arg));
  52             return FAILURE;
  53         } else {
  54             is_package_arg = false;
  55         }
  56     }
  57
  58     let matches = match opts.parse(env::args().skip(1).take_while(|a| a != "--")) {
  59         Ok(m) => m,
  60         Err(e) => {
  61             print_usage_to_stderr(&opts, &e.to_string());
  62             return FAILURE;
  63         }
  64     };
  65
  66     let verbosity = match (matches.opt_present("v"), matches.opt_present("q")) {
  67         (false, false) => Verbosity::Normal,
  68         (false, true) => Verbosity::Quiet,
  69         (true, false) => Verbosity::Verbose,
  70         (true, true) => {
  71             print_usage_to_stderr(&opts, "quiet mode and verbose mode are not compatible");
  72             return FAILURE;
  73         }
  74     };
  75
  76     if matches.opt_present("h") {
  77         print_usage_to_stdout(&opts, "");
  78         return SUCCESS;
  79     }
  80
  81     if matches.opt_present("version") {
  82         return handle_command_status(get_version(), &opts);
  83     }
  84
  85     let strategy = CargoFmtStrategy::from_matches(&matches);
  86     handle_command_status(format_crate(verbosity, &strategy), &opts)
  87 }
  88
  89 macro_rules! print_usage {
  90     ($print:ident, $opts:ident, $reason:expr) => {{
  91         let msg = format!("{}\nusage: cargo fmt [options]", $reason);
  92         $print!(
  93             "{}\nThis utility formats all bin and lib files of the current crate using rustfmt. \
  94              Arguments after `--` are passed to rustfmt.",
  95             $opts.usage(&msg)
  96         );
  97     }};
  98 }
  99
 100 fn print_usage_to_stdout(opts: &Options, reason: &str) {
 101     print_usage!(println, opts, reason);
 102 }
 103
 104 fn print_usage_to_stderr(opts: &Options, reason: &str) {
 105     print_usage!(eprintln, opts, reason);
 106 }
 107
 108 #[derive(Debug, Clone, Copy, PartialEq)]
 109 pub enum Verbosity {
 110     Verbose,
 111     Normal,
 112     Quiet,
 113 }
 114
 115 fn handle_command_status(status: Result<i32, io::Error>, opts: &getopts::Options) -> i32 {
 116     match status {
 117         Err(e) => {
 118             print_usage_to_stderr(opts, &e.to_string());
 119             FAILURE
 120         }
 121         Ok(status) => status,
 122     }
 123 }
 124
 125 fn get_version() -> Result<i32, io::Error> {
 126     let mut command = Command::new("rustfmt")
 127         .stdout(std::process::Stdio::inherit())
 128         .args(&[String::from("--version")])
 129         .spawn()
 130         .map_err(|e| match e.kind() {
 131             io::ErrorKind::NotFound => io::Error::new(
 132                 io::ErrorKind::Other,
 133                 "Could not run rustfmt, please make sure it is in your PATH.",
 134             ),
 135             _ => e,
 136         })?;
 137     let result = command.wait()?;
 138     if result.success() {
 139         Ok(SUCCESS)
 140     } else {
 141         Ok(result.code().unwrap_or(SUCCESS))
 142     }
 143 }
 144
 145 fn format_crate(verbosity: Verbosity, strategy: &CargoFmtStrategy) -> Result<i32, io::Error> {
 146     let rustfmt_args = get_fmt_args();
 147     let targets = if rustfmt_args
 148         .iter()
 149         .any(|s| ["--print-config", "-h", "--help", "-V", "--version"].contains(&s.as_str()))
 150     {
 151         BTreeSet::new()
 152     } else {
 153         get_targets(strategy)?
 154     };
 155
 156     // Currently only bin and lib files get formatted.
 157     run_rustfmt(&targets, &rustfmt_args, verbosity)
 158 }
 159
 160 fn get_fmt_args() -> Vec<String> {
 161     // All arguments after -- are passed to rustfmt.
 162     env::args().skip_while(|a| a != "--").skip(1).collect()
 163 }
 164
 165 /// Target uses a `path` field for equality and hashing.
 166 #[derive(Debug)]
 167 pub struct Target {
 168     /// A path to the main source file of the target.
 169     path: PathBuf,
 170     /// A kind of target (e.g., lib, bin, example, ...).
 171     kind: String,
 172     /// Rust edition for this target.
 173     edition: String,
 174 }
 175
 176 impl Target {
 177     pub fn from_target(target: &cargo_metadata::Target) -> Self {
 178         let path = PathBuf::from(&target.src_path);
 179         let canonicalized = fs::canonicalize(&path).unwrap_or(path);
 180
 181         Target {
 182             path: canonicalized,
 183             kind: target.kind[0].clone(),
 184             edition: target.edition.clone(),
 185         }
 186     }
 187 }
 188
 189 impl PartialEq for Target {
 190     fn eq(&self, other: &Target) -> bool {
 191         self.path == other.path
 192     }
 193 }
 194
 195 impl PartialOrd for Target {
 196     fn partial_cmp(&self, other: &Target) -> Option<Ordering> {
 197         Some(self.path.cmp(&other.path))
 198     }
 199 }
 200
 201 impl Ord for Target {
 202     fn cmp(&self, other: &Target) -> Ordering {
 203         self.path.cmp(&other.path)
 204     }
 205 }
 206
 207 impl Eq for Target {}
 208
 209 impl Hash for Target {
 210     fn hash<H: Hasher>(&self, state: &mut H) {
 211         self.path.hash(state);
 212     }
 213 }
 214
 215 #[derive(Debug, PartialEq, Eq)]
 216 pub enum CargoFmtStrategy {
 217     /// Format every packages and dependencies.
 218     All,
 219     /// Format packages that are specified by the command line argument.
 220     Some(Vec<String>),
 221     /// Format the root packages only.
 222     Root,
 223 }
 224
 225 impl CargoFmtStrategy {
 226     pub fn from_matches(matches: &Matches) -> CargoFmtStrategy {
 227         match (matches.opt_present("all"), matches.opt_present("p")) {
 228             (false, false) => CargoFmtStrategy::Root,
 229             (true, _) => CargoFmtStrategy::All,
 230             (false, true) => CargoFmtStrategy::Some(matches.opt_strs("p")),
 231         }
 232     }
 233 }
 234
 235 /// Based on the specified `CargoFmtStrategy`, returns a set of main source files.
 236 fn get_targets(strategy: &CargoFmtStrategy) -> Result<BTreeSet<Target>, io::Error> {
 237     let mut targets = BTreeSet::new();
 238
 239     match *strategy {
 240         CargoFmtStrategy::Root => get_targets_root_only(&mut targets)?,
 241         CargoFmtStrategy::All => get_targets_recursive(None, &mut targets, &mut BTreeSet::new())?,
 242         CargoFmtStrategy::Some(ref hitlist) => get_targets_with_hitlist(hitlist, &mut targets)?,
 243     }
 244
 245     if targets.is_empty() {
 246         Err(io::Error::new(
 247             io::ErrorKind::Other,
 248             "Failed to find targets".to_owned(),
 249         ))
 250     } else {
 251         Ok(targets)
 252     }
 253 }
 254
 255 fn get_targets_root_only(targets: &mut BTreeSet<Target>) -> Result<(), io::Error> {
 256     let metadata = get_cargo_metadata(None)?;
 257     let current_dir = env::current_dir()?.canonicalize()?;
 258     let current_dir_manifest = current_dir.join("Cargo.toml");
 259     let workspace_root_path = PathBuf::from(&metadata.workspace_root).canonicalize()?;
 260     let in_workspace_root = workspace_root_path == current_dir;
 261
 262     for package in metadata.packages {
 263         if in_workspace_root || PathBuf::from(&package.manifest_path) == current_dir_manifest {
 264             for target in package.targets {
 265                 targets.insert(Target::from_target(&target));
 266             }
 267         }
 268     }
 269
 270     Ok(())
 271 }
 272
 273 fn get_targets_recursive(
 274     manifest_path: Option<&Path>,
 275     mut targets: &mut BTreeSet<Target>,
 276     visited: &mut BTreeSet<String>,
 277 ) -> Result<(), io::Error> {
 278     let metadata = get_cargo_metadata(manifest_path)?;
 279
 280     for package in metadata.packages {
 281         add_targets(&package.targets, &mut targets);
 282
 283         // Look for local dependencies.
 284         for dependency in package.dependencies {
 285             if dependency.source.is_some() || visited.contains(&dependency.name) {
 286                 continue;
 287             }
 288
 289             let mut manifest_path = PathBuf::from(&package.manifest_path);
 290
 291             manifest_path.pop();
 292             manifest_path.push(&dependency.name);
 293             manifest_path.push("Cargo.toml");
 294
 295             if manifest_path.exists() {
 296                 visited.insert(dependency.name);
 297                 get_targets_recursive(Some(&manifest_path), &mut targets, visited)?;
 298             }
 299         }
 300     }
 301
 302     Ok(())
 303 }
 304
 305 fn get_targets_with_hitlist(
 306     hitlist: &[String],
 307     targets: &mut BTreeSet<Target>,
 308 ) -> Result<(), io::Error> {
 309     let metadata = get_cargo_metadata(None)?;
 310
 311     let mut workspace_hitlist: BTreeSet<&String> = BTreeSet::from_iter(hitlist);
 312
 313     for package in metadata.packages {
 314         if workspace_hitlist.remove(&package.name) {
 315             for target in package.targets {
 316                 targets.insert(Target::from_target(&target));
 317             }
 318         }
 319     }
 320
 321     if workspace_hitlist.is_empty() {
 322         Ok(())
 323     } else {
 324         let package = workspace_hitlist.iter().next().unwrap();
 325         Err(io::Error::new(
 326             io::ErrorKind::InvalidInput,
 327             format!("package `{}` is not a member of the workspace", package),
 328         ))
 329     }
 330 }
 331
 332 fn add_targets(target_paths: &[cargo_metadata::Target], targets: &mut BTreeSet<Target>) {
 333     for target in target_paths {
 334         targets.insert(Target::from_target(target));
 335     }
 336 }
 337
 338 fn run_rustfmt(
 339     targets: &BTreeSet<Target>,
 340     fmt_args: &[String],
 341     verbosity: Verbosity,
 342 ) -> Result<i32, io::Error> {
 343     let by_edition = targets
 344         .iter()
 345         .inspect(|t| {
 346             if verbosity == Verbosity::Verbose {
 347                 println!("[{} ({})] {:?}", t.kind, t.edition, t.path)
 348             }
 349         })
 350         .fold(BTreeMap::new(), |mut h, t| {
 351             h.entry(&t.edition).or_insert_with(Vec::new).push(&t.path);
 352             h
 353         });
 354
 355     let mut status = vec![];
 356     for (edition, files) in by_edition {
 357         let stdout = if verbosity == Verbosity::Quiet {
 358             std::process::Stdio::null()
 359         } else {
 360             std::process::Stdio::inherit()
 361         };
 362
 363         if verbosity == Verbosity::Verbose {
 364             print!("rustfmt");
 365             print!(" --edition {}", edition);
 366             fmt_args.iter().for_each(|f| print!(" {}", f));
 367             files.iter().for_each(|f| print!(" {}", f.display()));
 368             println!();
 369         }
 370
 371         let mut command = Command::new("rustfmt")
 372             .stdout(stdout)
 373             .args(files)
 374             .args(&["--edition", edition])
 375             .args(fmt_args)
 376             .spawn()
 377             .map_err(|e| match e.kind() {
 378                 io::ErrorKind::NotFound => io::Error::new(
 379                     io::ErrorKind::Other,
 380                     "Could not run rustfmt, please make sure it is in your PATH.",
 381                 ),
 382                 _ => e,
 383             })?;
 384
 385         status.push(command.wait()?);
 386     }
 387
 388     Ok(status
 389         .iter()
 390         .filter_map(|s| if s.success() { None } else { s.code() })
 391         .next()
 392         .unwrap_or(SUCCESS))
 393 }
 394
 395 fn get_cargo_metadata(manifest_path: Option<&Path>) -> Result<cargo_metadata::Metadata, io::Error> {
 396     let mut cmd = cargo_metadata::MetadataCommand::new();
 397     cmd.no_deps();
 398     if let Some(manifest_path) = manifest_path {
 399         cmd.manifest_path(manifest_path);
 400     }
 401     match cmd.exec() {
 402         Ok(metadata) => Ok(metadata),
 403         Err(error) => Err(io::Error::new(io::ErrorKind::Other, error.to_string())),
 404     }
 405 }