Merge pull request #10 from akoken/bug-fix-color-issue

Fix colorisation issue

src/main.rs

@@ -48,28 +48,38 @@ impl PartialEq for SearchResult {
 }
 
 pub fn search(query: &str, contents: &str, ignore_case: bool) -> Vec<SearchResult> {
+    let query_to_search = if ignore_case {
+        query.to_lowercase()
+    } else {
+        query.to_string()
+    };
+
     contents
         .lines()
         .enumerate()
         .filter_map(|(index, line)| {
-            let matches = if ignore_case {
-                line.to_lowercase().contains(&query.to_lowercase())
+            let line_to_search = if ignore_case {
+                line.to_lowercase()
             } else {
-                line.contains(query)
+                line.to_string()
             };
 
-            if matches {
-                let colored_line = if ignore_case {
-                    let mut colored_line = line.to_string();
-                    for (start, part) in line.to_lowercase().match_indices(&query.to_lowercase()) {
-                        let end = start + part.len();
-                        let colored_part = &line[start..end].red().to_string();
-                        colored_line.replace_range(start..end, colored_part);
-                    }
-                    colored_line
-                } else {
-                    line.replace(query, &query.red().to_string())
-                };
+            if line_to_search.contains(&query_to_search) {
+                let mut colored_line = String::new();
+                let mut last_match_end = 0;
+
+                let matches = line_to_search.match_indices(&query_to_search);
+
+                for (start, _) in matches {
+                    colored_line.push_str(&line[last_match_end..start]);
+
+                    let original_text = &line[start..start + query.len()];
+                    colored_line.push_str(&original_text.red().to_string());
+
+                    last_match_end = start + query.len();
+                }
+
+                colored_line.push_str(&line[last_match_end..]);
 
                 Some(SearchResult {
                     line_number: (index + 1) as u32,
@@ -81,45 +91,73 @@ pub fn search(query: &str, contents: &str, ignore_case: bool) -> Vec<SearchResul
         })
         .collect()
 }
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
+    fn strip_color_codes(s: &str) -> String {
+        // Use regex or replace ANSI codes directly if needed
+        s.replace("\x1b[31m", "").replace("\x1b[0m", "")
+    }
+
+    #[test]
+    fn test_case_sensitive_colorization() {
+        let query = "Warning";
+        let contents = "Warning: This is a test.\nAnother line with Warning.\nFinal line.";
+        let ignore_case = false;
+
+        let results = search(query, contents, ignore_case);
+
+        assert_eq!(results.len(), 2);
+
+        let expected_line1 = format!("{}: This is a test.", "Warning".red());
+        let expected_line2 = format!("Another line with {}.", "Warning".red());
+
+        assert_eq!(results[0].line_text, expected_line1);
+        assert_eq!(results[1].line_text, expected_line2);
+    }
+
+    #[test]
+    fn test_case_insensitive_colorization() {
+        let query = "warning";
+        let contents = "Warning: This is a test.\nAnother line with warning.\nFinal line.";
+        let ignore_case = true;
+
+        let results = search(query, contents, ignore_case);
+
+        assert_eq!(results.len(), 2);
+
+        let expected_line1 = format!("{}: This is a test.", "Warning".red());
+        let expected_line2 = format!("Another line with {}.", "warning".red());
+
+        assert_eq!(results[0].line_text, expected_line1);
+        assert_eq!(results[1].line_text, expected_line2);
+    }
+
     #[test]
-    fn case_sensitive() {
-        let query = "duct";
-        let contents = "\
-Rust:
-safe, fast, productive.
-Pick three.
-Duct tape.";
-
-        let expected = vec![SearchResult {
-            line_number: 2,
-            line_text: "safe, fast, pro".to_string() + &"duct".red().to_string() + "ive.",
-        }];
-        assert_eq!(expected, search(query, contents, false));
+    fn test_no_colorization_in_non_matching_line() {
+        let query = "Error";
+        let contents = "Warning: This is a test.\nAnother line with Warning.\nFinal line.";
+        let ignore_case = true;
+
+        let results = search(query, contents, ignore_case);
+
+        assert_eq!(results.len(), 0);
     }
 
     #[test]
-    fn case_insensitive() {
-        let query = "rUsT";
-        let contents = "\
-Rust:
-safe, fast, productive.
-Pick three.
-Trust me.";
-
-        let expected = vec![
-            SearchResult {
-                line_number: 1,
-                line_text: "Rust".red().to_string() + ":",
-            },
-            SearchResult {
-                line_number: 4,
-                line_text: "T".to_string() + &"rust".red().to_string() + " me.",
-            },
-        ];
-        assert_eq!(expected, search(query, contents, true));
+    fn test_colorization_strip() {
+        let query = "Warning";
+        let contents = "Warning: Check the warning levels.\nIgnore warnings on final line.";
+        let ignore_case = false;
+
+        let results = search(query, contents, ignore_case);
+
+        for result in results {
+            let stripped_text = strip_color_codes(&result.line_text);
+            assert!(stripped_text.contains(query));
+            assert!(!stripped_text.contains("\x1b["));
+        }
     }
 }