diff --git a/src/lib.rs b/src/lib.rs index 4b58e14f..96244abd 100644 --- a/src/lib.rs +++ b/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")] diff --git a/src/reader.rs b/src/reader.rs index e6a8a103..1f5d8b62 100644 --- a/src/reader.rs +++ b/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, + 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(f: R, timeout: Option) -> NBReader { + /// - options: see `Options` + pub fn new(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,7 +397,7 @@ 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()), @@ -384,10 +405,44 @@ mod tests { ); } + #[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"); diff --git a/src/session.rs b/src/session.rs index 467a5b9f..7bfef797 100644 --- a/src/session.rs +++ b/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 { } impl StreamSession { - pub fn new(reader: R, writer: W, timeout_ms: Option) -> Self { + pub fn new(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) -> Result { + fn new(process: PtyProcess, options: Options) -> Result { 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) -> Result) -> Result { + spawn_with_options( + command, + Options { + timeout_ms, + strip_ansi_escape_codes: false, + }, + ) +} + +/// See `spawn` +pub fn spawn_with_options(command: Command, options: Options) -> Result { #[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( writer: W, timeout_ms: Option, ) -> StreamSession { - StreamSession::new(reader, writer, timeout_ms) + StreamSession::new( + reader, + writer, + Options { + timeout_ms, + strip_ansi_escape_codes: false, + }, + ) } #[cfg(test)]