diff --git a/src/executor.rs b/src/executor.rs index e9ef47de..87b320dc 100644 --- a/src/executor.rs +++ b/src/executor.rs @@ -1,6 +1,6 @@ //! Utilities for command execution use crate::error::TopgradeError; -use crate::utils::Check; +use crate::utils::{Check, CheckWithCodes}; use anyhow::Result; use log::{debug, trace}; use std::ffi::{OsStr, OsString}; @@ -166,6 +166,13 @@ impl Executor { pub fn check_run(&mut self) -> Result<()> { self.spawn()?.wait()?.check() } + + /// An extension of `check_run` that allows you to set a sequence of codes + /// that can indicate success of a script + #[allow(dead_code)] + pub fn check_run_with_codes(&mut self, codes: &[i32]) -> Result<()> { + self.spawn()?.wait()?.check_with_codes(codes) + } } pub enum ExecutorOutput { @@ -232,6 +239,15 @@ impl Check for ExecutorExitStatus { } } +impl CheckWithCodes for ExecutorExitStatus { + fn check_with_codes(self, codes: &[i32]) -> Result<()> { + match self { + ExecutorExitStatus::Wet(e) => e.check_with_codes(codes), + ExecutorExitStatus::Dry => Ok(()), + } + } +} + /// Extension methods for `std::process::Command` pub trait CommandExt { /// Run the command, wait for it to complete, check the return code and decode the output as UTF-8. diff --git a/src/steps/zsh.rs b/src/steps/zsh.rs index 85cc3c52..e769d836 100644 --- a/src/steps/zsh.rs +++ b/src/steps/zsh.rs @@ -142,5 +142,5 @@ pub fn run_oh_my_zsh(ctx: &ExecutionContext) -> Result<()> { .execute("zsh") .env("ZSH", &oh_my_zsh) .arg(&oh_my_zsh.join("tools/upgrade.sh")) - .check_run() + .check_run_with_codes(&[80]) } diff --git a/src/utils.rs b/src/utils.rs index 20ef63c6..7442f69b 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -28,6 +28,22 @@ impl Check for Output { } } +pub trait CheckWithCodes { + fn check_with_codes(self, codes: &[i32]) -> Result<()>; +} + +impl CheckWithCodes for ExitStatus { + fn check_with_codes(self, codes: &[i32]) -> Result<()> { + // Set the default to be -1 because the option represents a signal termination + let code = self.code().unwrap_or(-1); + if self.success() || codes.contains(&code) { + Ok(()) + } else { + Err(TopgradeError::ProcessFailed(self).into()) + } + } +} + pub trait PathExt where Self: Sized,