Fixes issue with ohmyzsh returning 80 on a successful run (#570)

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.

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])
 }

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,