Add option to skip ansi escape codes

ansi escape codes are just the byte 27 inserted in a string,
followed by "[", zero or more numbers separated by ";" and ending
with a letter.
We optionally just skip those in our non-blocking reader in order
to efficently ignore them during matching.

Signed-off-by: phaer <[email protected]>

Author: phaer

README.md

@@ -38,7 +38,7 @@ use rexpect::spawn;
 use rexpect::error::*;
 
 fn do_ftp() -> Result<(), Error> {
-    let mut p = spawn("ftp speedtest.tele2.net", Some(30_000))?;
+    let mut p = spawn("ftp speedtest.tele2.net", Some(30_000), false)?;
     p.exp_regex("Name \\(.*\\):")?;
     p.send_line("anonymous")?;
     p.exp_string("Password")?;
@@ -65,7 +65,7 @@ use rexpect::spawn_bash;
 use rexpect::error::*;
 
 fn do_bash() -> Result<(), Error> {
-    let mut p = spawn_bash(Some(2000))?;
+    let mut p = spawn_bash(Some(2000), false)?;
 
     // case 1: wait until program is done
     p.send_line("hostname")?;
@@ -115,7 +115,7 @@ use rexpect::spawn_bash;
 use rexpect::error::*;
 
 fn do_bash_jobcontrol() -> Result<(), Error> {
-    let mut p = spawn_bash(Some(1000))?;
+    let mut p = spawn_bash(Some(1000), false)?;
     p.execute("ping 8.8.8.8", "bytes of data")?;
     p.send_control('z')?;
     p.wait_for_prompt()?;

examples/bash.rs

@@ -2,7 +2,7 @@ use rexpect::error::Error;
 use rexpect::spawn_bash;
 
 fn main() -> Result<(), Error> {
-    let mut p = spawn_bash(Some(1000))?;
+    let mut p = spawn_bash(Some(1000), false)?;
     p.execute("ping 8.8.8.8", "bytes")?;
     p.send_control('z')?;
     p.wait_for_prompt()?;

examples/bash_read.rs

@@ -2,7 +2,7 @@ use rexpect::error::Error;
 use rexpect::spawn_bash;
 
 fn main() -> Result<(), Error> {
-    let mut p = spawn_bash(Some(2000))?;
+    let mut p = spawn_bash(Some(2000), false)?;
 
     // case 1: wait until program is done
     p.send_line("hostname")?;

examples/exit_code.rs

@@ -7,13 +7,13 @@ use rexpect::spawn;
 /// cat exited with code 1
 /// Output (stdout and stderr): cat: /this/does/not/exist: No such file or directory
 fn main() -> Result<(), Error> {
-    let p = spawn("cat /etc/passwd", Some(2000))?;
+    let p = spawn("cat /etc/passwd", Some(2000), false)?;
     match p.process.wait() {
         Ok(wait::WaitStatus::Exited(_, 0)) => println!("cat exited with code 0, all good!"),
         _ => println!("cat exited with code >0, or it was killed"),
     }
 
-    let mut p = spawn("cat /this/does/not/exist", Some(2000))?;
+    let mut p = spawn("cat /this/does/not/exist", Some(2000), false)?;
     match p.process.wait() {
         Ok(wait::WaitStatus::Exited(_, 0)) => println!("cat succeeded"),
         Ok(wait::WaitStatus::Exited(_, c)) => {

examples/ftp.rs

@@ -2,7 +2,7 @@ use rexpect::error::Error;
 use rexpect::spawn;
 
 fn main() -> Result<(), Error> {
-    let mut p = spawn("ftp speedtest.tele2.net", Some(2000))?;
+    let mut p = spawn("ftp speedtest.tele2.net", Some(2000), false)?;
     p.exp_regex("Name \\(.*\\):")?;
     p.send_line("anonymous")?;
     p.exp_string("Password")?;

examples/repl.rs

@@ -12,7 +12,7 @@ fn ed_session() -> Result<PtyReplSession, Error> {
 
         // used for `wait_for_prompt()`
         prompt: "> ".to_string(),
-        pty_session: spawn("/bin/ed -p '> '", Some(2000))?,
+        pty_session: spawn("/bin/ed -p '> '", Some(2000), false)?,
         // command which is sent when the instance of this struct is dropped
         // in the below example this is not needed, but if you don't explicitly
         // exit a REPL then rexpect tries to send a SIGTERM and depending on the repl

src/lib.rs

@@ -20,7 +20,7 @@
 //! use rexpect::error::Error;
 //!
 //! fn main() -> Result<(), Error> {
-//!     let mut p = spawn("ftp speedtest.tele2.net", Some(2000))?;
+//!     let mut p = spawn("ftp speedtest.tele2.net", Some(2000), false)?;
 //!     p.exp_regex("Name \\(.*\\):")?;
 //!     p.send_line("anonymous")?;
 //!     p.exp_string("Password")?;
@@ -50,7 +50,7 @@
 //! use rexpect::error::Error;
 //!
 //! fn main() -> Result<(), Error> {
-//!     let mut p = spawn_bash(Some(30_000))?;
+//!     let mut p = spawn_bash(Some(30_000), false)?;
 //!     p.execute("ping 8.8.8.8", "bytes of data")?;
 //!     p.send_control('z')?;
 //!     p.wait_for_prompt()?;

src/reader.rs

@@ -119,13 +119,19 @@ impl NBReader {
     /// - timeout:
     ///  + `None`: read_until is blocking forever. This is probably not what you want
     ///  + `Some(millis)`: after millis milliseconds a timeout error is raised
-    pub fn new<R: Read + Send + 'static>(f: R, timeout: Option<u64>) -> NBReader {
+    /// - strip_ansi_escape_codes: Whether to filter out escape codes, such as colors.
+    pub fn new<R: Read + Send + 'static>(
+        f: R,
+        timeout: Option<u64>,
+        strip_ansi_escape_codes: bool
+    ) -> NBReader {
         let (tx, rx) = channel();
 
         // spawn a thread which reads one char and sends it to tx
         thread::spawn(move || -> Result<(), Error> {
             let mut reader = BufReader::new(f);
             let mut byte = [0u8];
+
             loop {
                 match reader.read(&mut byte) {
                     Ok(0) => {
@@ -134,12 +140,21 @@ impl NBReader {
                         break;
                     }
                     Ok(_) => {
-                        tx.send(Ok(PipedChar::Char(byte[0])))
-                            .map_err(|_| Error::MpscSendError)?;
+                        if strip_ansi_escape_codes && byte[0] == 27 {
+                            while let Ok(_) = reader.read(&mut byte) {
+                                if char::from(byte[0]).is_alphabetic() {
+                                    break;
+                                }
+                            }
+                        }
+                        else {
+                            tx.send(Ok(PipedChar::Char(byte[0])))
+                              .map_err(|_| Error::MpscSendError)?;
+                        }
                     }
                     Err(error) => {
                         tx.send(Err(PipeError::IO(error)))
-                            .map_err(|_| Error::MpscSendError)?;
+                          .map_err(|_| Error::MpscSendError)?;
                     }
                 }
             }
@@ -208,7 +223,7 @@ impl NBReader {
     /// // instead of a Cursor you would put your process output or file here
     /// let f = Cursor::new("Hello, miss!\n\
     ///                         What do you mean: 'miss'?");
-    /// let mut e = NBReader::new(f, None);
+    /// let mut e = NBReader::new(f, None, false);
     ///
     /// let (first_line, _) = e.read_until(&ReadUntil::String('\n'.to_string())).unwrap();
     /// assert_eq!("Hello, miss!", &first_line);
@@ -230,6 +245,7 @@ impl NBReader {
 
         loop {
             self.read_into_buffer()?;
+
             if let Some(tuple_pos) = find(needle, &self.buffer, self.eof) {
                 let first = self.buffer.drain(..tuple_pos.0).collect();
                 let second = self.buffer.drain(..tuple_pos.1 - tuple_pos.0).collect();
@@ -287,7 +303,7 @@ mod tests {
     #[test]
     fn test_expect_melon() {
         let f = io::Cursor::new("a melon\r\n");
-        let mut r = NBReader::new(f, None);
+        let mut r = NBReader::new(f, None, false);
         assert_eq!(
             ("a melon".to_string(), "\r\n".to_string()),
             r.read_until(&ReadUntil::String("\r\n".to_string()))
@@ -304,7 +320,7 @@ mod tests {
     #[test]
     fn test_regex() {
         let f = io::Cursor::new("2014-03-15");
-        let mut r = NBReader::new(f, None);
+        let mut r = NBReader::new(f, None, false);
         let re = Regex::new(r"^\d{4}-\d{2}-\d{2}$").unwrap();
         assert_eq!(
             ("".to_string(), "2014-03-15".to_string()),
@@ -316,7 +332,7 @@ mod tests {
     #[test]
     fn test_regex2() {
         let f = io::Cursor::new("2014-03-15");
-        let mut r = NBReader::new(f, None);
+        let mut r = NBReader::new(f, None, false);
         let re = Regex::new(r"-\d{2}-").unwrap();
         assert_eq!(
             ("2014".to_string(), "-03-".to_string()),
@@ -328,7 +344,7 @@ mod tests {
     #[test]
     fn test_nbytes() {
         let f = io::Cursor::new("abcdef");
-        let mut r = NBReader::new(f, None);
+        let mut r = NBReader::new(f, None, false);
         assert_eq!(
             ("".to_string(), "ab".to_string()),
             r.read_until(&ReadUntil::NBytes(2)).expect("2 bytes")
@@ -346,7 +362,7 @@ mod tests {
     #[test]
     fn test_any_with_multiple_possible_matches() {
         let f = io::Cursor::new("zero one two three four five");
-        let mut r = NBReader::new(f, None);
+        let mut r = NBReader::new(f, None, false);
 
         let result = r
             .read_until(&ReadUntil::Any(vec![
@@ -361,7 +377,7 @@ mod tests {
     #[test]
     fn test_any_with_same_start_different_length() {
         let f = io::Cursor::new("hi hello");
-        let mut r = NBReader::new(f, None);
+        let mut r = NBReader::new(f, None, false);
 
         let result = r
             .read_until(&ReadUntil::Any(vec![
@@ -376,18 +392,29 @@ mod tests {
     #[test]
     fn test_eof() {
         let f = io::Cursor::new("lorem ipsum dolor sit amet");
-        let mut r = NBReader::new(f, None);
+        let mut r = NBReader::new(f, None, false);
         r.read_until(&ReadUntil::NBytes(2)).expect("2 bytes");
         assert_eq!(
             ("".to_string(), "rem ipsum dolor sit amet".to_string()),
             r.read_until(&ReadUntil::EOF).expect("reading until EOF")
         );
     }
 
+    #[test]
+    fn test_skip_ansi_codes() {
+        let f = io::Cursor::new("\x1b[31;1;4mHello\x1b[0m");
+        let mut r = NBReader::new(f, None, true);
+        let bytes = r.read_until(&ReadUntil::String("Hello".to_string())).unwrap();
+        assert_eq!(bytes, ("".to_string(), "Hello".to_string()));
+        assert_eq!(None, r.try_read());
+
+    }
+
+
     #[test]
     fn test_try_read() {
         let f = io::Cursor::new("lorem");
-        let mut r = NBReader::new(f, None);
+        let mut r = NBReader::new(f, None, false);
         let bytes = r.read_until(&ReadUntil::NBytes(4)).unwrap();
         assert!(bytes.0.is_empty());
         assert_eq!(bytes.1, "lore");