From 400f93b0df875847709ac84da42f32ac7f2ab0e6 Mon Sep 17 00:00:00 2001
From: Krzysztof Pawlik <krzysiek.pawlik@people.pl>
Date: Thu, 4 Jul 2024 19:07:01 +0200
Subject: [PATCH] feat(character): allow specifying more exit codes as
 successful

---
 .github/config-schema.json |  9 +++++++++
 docs/config/README.md      |  1 +
 src/configs/character.rs   |  2 ++
 src/modules/character.rs   | 38 ++++++++++++++++++++++++++++++++++++--
 4 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/.github/config-schema.json b/.github/config-schema.json
index 5dc4696e6..dc2ac9e05 100644
--- a/.github/config-schema.json
+++ b/.github/config-schema.json
@@ -142,6 +142,7 @@
         "disabled": false,
         "error_symbol": "[❯](bold red)",
         "format": "$symbol ",
+        "success_exit_codes": [],
         "success_symbol": "[❯](bold green)",
         "vimcmd_replace_one_symbol": "[❮](bold purple)",
         "vimcmd_replace_symbol": "[❮](bold purple)",
@@ -2327,6 +2328,14 @@
           "default": "[❯](bold red)",
           "type": "string"
         },
+        "success_exit_codes": {
+          "default": [],
+          "type": "array",
+          "items": {
+            "type": "integer",
+            "format": "int32"
+          }
+        },
         "vimcmd_symbol": {
           "default": "[❮](bold green)",
           "type": "string"
diff --git a/docs/config/README.md b/docs/config/README.md
index b5ab46b53..19345a12a 100644
--- a/docs/config/README.md
+++ b/docs/config/README.md
@@ -741,6 +741,7 @@ are only supported in fish due to [upstream issues with mode detection in zsh](h
 | `format`                    | `'$symbol '`         | The format string used before the text input.                                           |
 | `success_symbol`            | `'[❯](bold green)'`  | The format string used before the text input if the previous command succeeded.         |
 | `error_symbol`              | `'[❯](bold red)'`    | The format string used before the text input if the previous command failed.            |
+| `success_exit_codes`        | `[]`                 | List of additional exit codes that are rendered as a success (0 is always a success).   |
 | `vimcmd_symbol`             | `'[❮](bold green)'`  | The format string used before the text input if the shell is in vim normal mode.        |
 | `vimcmd_replace_one_symbol` | `'[❮](bold purple)'` | The format string used before the text input if the shell is in vim `replace_one` mode. |
 | `vimcmd_replace_symbol`     | `'[❮](bold purple)'` | The format string used before the text input if the shell is in vim replace mode.       |
diff --git a/src/configs/character.rs b/src/configs/character.rs
index c6719a671..f235accab 100644
--- a/src/configs/character.rs
+++ b/src/configs/character.rs
@@ -11,6 +11,7 @@ pub struct CharacterConfig<'a> {
     pub format: &'a str,
     pub success_symbol: &'a str,
     pub error_symbol: &'a str,
+    pub success_exit_codes: Vec<i32>,
     #[serde(alias = "vicmd_symbol")]
     pub vimcmd_symbol: &'a str,
     pub vimcmd_visual_symbol: &'a str,
@@ -25,6 +26,7 @@ impl<'a> Default for CharacterConfig<'a> {
             format: "$symbol ",
             success_symbol: "[❯](bold green)",
             error_symbol: "[❯](bold red)",
+            success_exit_codes: vec![],
             vimcmd_symbol: "[❮](bold green)",
             vimcmd_visual_symbol: "[❮](bold yellow)",
             vimcmd_replace_symbol: "[❮](bold purple)",
diff --git a/src/modules/character.rs b/src/modules/character.rs
index 232b38288..341b2fb19 100644
--- a/src/modules/character.rs
+++ b/src/modules/character.rs
@@ -25,9 +25,14 @@ pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
     let config: CharacterConfig = CharacterConfig::try_load(module.config);
 
     let props = &context.properties;
-    let exit_code = props.status_code.as_deref().unwrap_or("0");
+    let exit_code = props
+        .status_code
+        .as_deref()
+        .unwrap_or("0")
+        .parse::<i32>()
+        .unwrap_or(0);
     let keymap = props.keymap.as_str();
-    let exit_success = exit_code == "0";
+    let exit_success = exit_code == 0 || config.success_exit_codes.contains(&exit_code);
 
     // Match shell "keymap" names to normalized vi modes
     // NOTE: in vi mode, fish reports normal mode as "default".
@@ -109,6 +114,35 @@ mod test {
         }
     }
 
+    #[test]
+    fn success_exit_codes_handling() {
+        let module_config = toml::toml! {
+            [character]
+            success_exit_codes = [146, 42]
+        };
+        let expected_success = Some(format!("{} ", Color::Green.bold().paint("❯")));
+        let expected_fail = Some(format!("{} ", Color::Red.bold().paint("❯")));
+
+        let exit_values_success = [0, 42, 146];
+        let exit_values_fail = [-5000, 1, 145, 147];
+
+        for status in &exit_values_success {
+            let actual = ModuleRenderer::new("character")
+                .config(module_config.clone())
+                .status(*status)
+                .collect();
+            assert_eq!(expected_success, actual);
+        }
+
+        for status in &exit_values_fail {
+            let actual = ModuleRenderer::new("character")
+                .config(module_config.clone())
+                .status(*status)
+                .collect();
+            assert_eq!(expected_fail, actual);
+        }
+    }
+
     #[test]
     fn custom_symbol() {
         let expected_fail = Some(format!("{} ", Color::Red.bold().paint("✖")));