Add option to skip ansi escape codes

src/lib.rs

@@ -70,7 +70,7 @@ pub mod reader;
 pub mod session;
 
 pub use reader::ReadUntil;
-pub use session::{spawn, spawn_bash, spawn_python, spawn_stream};
+pub use session::{spawn, spawn_bash, spawn_python, spawn_stream, spawn_with_options};
 
 // include the README.md here to test its doc
 #[doc = include_str!("../README.md")]

src/reader.rs

@@ -98,6 +98,18 @@ pub fn find(needle: &ReadUntil, buffer: &str, eof: bool) -> Option<(usize, usize
     }
 }
 
+/// Options for 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
+/// - strip_ansi_escape_codes: Whether to filter out escape codes, such as colors.
+#[derive(Default)]
+pub struct Options {
+    pub timeout_ms: Option<u64>,
+    pub strip_ansi_escape_codes: bool,
+}
+
 /// Non blocking reader
 ///
 /// Typically you'd need that to check for output of a process without blocking your thread.
@@ -116,16 +128,16 @@ impl NBReader {
     /// # Arguments:
     ///
     /// - f: file like object
-    /// - 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 {
+    /// - options: see `Options`
+    pub fn new<R: Read + Send + 'static>(f: R, options: Options) -> 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];
+            let mut in_escape_code = false;
+
             loop {
                 match reader.read(&mut byte) {
                     Ok(0) => {
@@ -134,8 +146,16 @@ impl NBReader {
                         break;
                     }
                     Ok(_) => {
-                        tx.send(Ok(PipedChar::Char(byte[0])))
-                            .map_err(|_| Error::MpscSendError)?;
+                        if options.strip_ansi_escape_codes && byte[0] == 27 {
+                            in_escape_code = true;
+                        } else if options.strip_ansi_escape_codes && in_escape_code {
+                            if char::from(byte[0]).is_alphabetic() {
+                                in_escape_code = false;
+                            }
+                        } else {
+                            tx.send(Ok(PipedChar::Char(byte[0])))
+                                .map_err(|_| Error::MpscSendError)?;
+                        }
                     }
                     Err(error) => {
                         tx.send(Err(PipeError::IO(error)))
@@ -153,7 +173,7 @@ impl NBReader {
             reader: rx,
             buffer: String::with_capacity(1024),
             eof: false,
-            timeout: timeout.map(time::Duration::from_millis),
+            timeout: options.timeout_ms.map(time::Duration::from_millis),
         }
     }
 
@@ -204,11 +224,11 @@ impl NBReader {
     ///
     /// ```
     /// # use std::io::Cursor;
-    /// use rexpect::reader::{NBReader, ReadUntil, Regex};
+    /// use rexpect::reader::{NBReader, ReadUntil, Regex, Options};
     /// // 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, Options::default());
     ///
     /// let (first_line, _) = e.read_until(&ReadUntil::String('\n'.to_string())).unwrap();
     /// assert_eq!("Hello, miss!", &first_line);
@@ -230,6 +250,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 +308,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, Options::default());
         assert_eq!(
             ("a melon".to_string(), "\r\n".to_string()),
             r.read_until(&ReadUntil::String("\r\n".to_string()))
@@ -304,7 +325,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, Options::default());
         let re = Regex::new(r"^\d{4}-\d{2}-\d{2}$").unwrap();
         assert_eq!(
             ("".to_string(), "2014-03-15".to_string()),
@@ -316,7 +337,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, Options::default());
         let re = Regex::new(r"-\d{2}-").unwrap();
         assert_eq!(
             ("2014".to_string(), "-03-".to_string()),
@@ -328,7 +349,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, Options::default());
         assert_eq!(
             ("".to_string(), "ab".to_string()),
             r.read_until(&ReadUntil::NBytes(2)).expect("2 bytes")
@@ -346,7 +367,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, Options::default());
 
         let result = r
             .read_until(&ReadUntil::Any(vec![
@@ -361,7 +382,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, Options::default());
 
         let result = r
             .read_until(&ReadUntil::Any(vec![
@@ -376,18 +397,52 @@ 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, Options::default());
         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_partial_ansi_code() {
+        let f = io::Cursor::new("\x1b[31;1;4mHello\x1b[1");
+        let mut r = NBReader::new(
+            f,
+            Options {
+                timeout_ms: None,
+                strip_ansi_escape_codes: 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_skip_ansi_codes() {
+        let f = io::Cursor::new("\x1b[31;1;4mHello\x1b[0m");
+        let mut r = NBReader::new(
+            f,
+            Options {
+                timeout_ms: None,
+                strip_ansi_escape_codes: 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, Options::default());
         let bytes = r.read_until(&ReadUntil::NBytes(4)).unwrap();
         assert!(bytes.0.is_empty());
         assert_eq!(bytes.1, "lore");

src/session.rs

@@ -2,8 +2,8 @@
 
 use crate::error::Error; // load error-chain
 use crate::process::PtyProcess;
-pub use crate::reader::ReadUntil;
 use crate::reader::{NBReader, Regex};
+pub use crate::reader::{Options, ReadUntil};
 use std::fs::File;
 use std::io::prelude::*;
 use std::io::LineWriter;
@@ -17,10 +17,10 @@ pub struct StreamSession<W: Write> {
 }
 
 impl<W: Write> StreamSession<W> {
-    pub fn new<R: Read + Send + 'static>(reader: R, writer: W, timeout_ms: Option<u64>) -> Self {
+    pub fn new<R: Read + Send + 'static>(reader: R, writer: W, options: Options) -> Self {
         Self {
             writer: LineWriter::new(writer),
-            reader: NBReader::new(reader, timeout_ms),
+            reader: NBReader::new(reader, options),
         }
     }
 
@@ -172,11 +172,29 @@ impl DerefMut for PtySession {
 }
 
 /// Start a process in a tty session, write and read from it
+///
+/// # Example
+///
+/// ```
+///
+/// use rexpect::spawn;
+/// # use rexpect::error::Error;
+///
+/// # fn main() {
+///     # || -> Result<(), Error> {
+/// let mut s = spawn("cat", Some(1000))?;
+/// s.send_line("hello, polly!")?;
+/// let line = s.read_line()?;
+/// assert_eq!("hello, polly!", line);
+///         # Ok(())
+///     # }().expect("test failed");
+/// # }
+/// ```
 impl PtySession {
-    fn new(process: PtyProcess, timeout_ms: Option<u64>) -> Result<Self, Error> {
+    fn new(process: PtyProcess, options: Options) -> Result<Self, Error> {
         let f = process.get_file_handle()?;
         let reader = f.try_clone()?;
-        let stream = StreamSession::new(reader, f, timeout_ms);
+        let stream = StreamSession::new(reader, f, options);
         Ok(Self { process, stream })
     }
 }
@@ -214,14 +232,25 @@ pub fn spawn(program: &str, timeout_ms: Option<u64>) -> Result<PtySession, Error
 
 /// See `spawn`
 pub fn spawn_command(command: Command, timeout_ms: Option<u64>) -> Result<PtySession, Error> {
+    spawn_with_options(
+        command,
+        Options {
+            timeout_ms,
+            strip_ansi_escape_codes: false,
+        },
+    )
+}
+
+/// See `spawn`
+pub fn spawn_with_options(command: Command, options: Options) -> Result<PtySession, Error> {
     #[cfg(feature = "which")]
     {
         let _ = which::which(command.get_program())?;
     }
     let mut process = PtyProcess::new(command)?;
-    process.set_kill_timeout(timeout_ms);
+    process.set_kill_timeout(options.timeout_ms);
 
-    PtySession::new(process, timeout_ms)
+    PtySession::new(process, options)
 }
 
 /// A repl session: e.g. bash or the python shell:
@@ -407,7 +436,14 @@ pub fn spawn_stream<R: Read + Send + 'static, W: Write>(
     writer: W,
     timeout_ms: Option<u64>,
 ) -> StreamSession<W> {
-    StreamSession::new(reader, writer, timeout_ms)
+    StreamSession::new(
+        reader,
+        writer,
+        Options {
+            timeout_ms,
+            strip_ansi_escape_codes: false,
+        },
+    )
 }
 
 #[cfg(test)]