From 6b37f7e8c8246d0b1fa3fb236229eee95dea16fe Mon Sep 17 00:00:00 2001 From: phaer Date: Sun, 23 Apr 2023 21:56:39 +0200 Subject: [PATCH 1/4] 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 --- README.md | 6 ++-- examples/bash.rs | 2 +- examples/bash_read.rs | 2 +- examples/exit_code.rs | 4 +-- examples/ftp.rs | 2 +- examples/repl.rs | 2 +- src/lib.rs | 4 +-- src/reader.rs | 53 ++++++++++++++++++++++++--------- src/session.rs | 68 +++++++++++++++++++++++++++---------------- 9 files changed, 94 insertions(+), 49 deletions(-) diff --git a/README.md b/README.md index 0bac2320..269b168e 100644 --- a/README.md +++ b/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()?; diff --git a/examples/bash.rs b/examples/bash.rs index a95ac46b..57e7677e 100644 --- a/examples/bash.rs +++ b/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()?; diff --git a/examples/bash_read.rs b/examples/bash_read.rs index 4522a062..fb3bf4c6 100644 --- a/examples/bash_read.rs +++ b/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")?; diff --git a/examples/exit_code.rs b/examples/exit_code.rs index 0964f6f8..9c52bacb 100644 --- a/examples/exit_code.rs +++ b/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)) => { diff --git a/examples/ftp.rs b/examples/ftp.rs index d1571879..0ace5dd8 100644 --- a/examples/ftp.rs +++ b/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")?; diff --git a/examples/repl.rs b/examples/repl.rs index 5c0c257f..54e2e40d 100644 --- a/examples/repl.rs +++ b/examples/repl.rs @@ -12,7 +12,7 @@ fn ed_session() -> Result { // 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 diff --git a/src/lib.rs b/src/lib.rs index 4b58e14f..9e9e1a82 100644 --- a/src/lib.rs +++ b/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()?; diff --git a/src/reader.rs b/src/reader.rs index e6a8a103..d013584c 100644 --- a/src/reader.rs +++ b/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(f: R, timeout: Option) -> NBReader { + /// - strip_ansi_escape_codes: Whether to filter out escape codes, such as colors. + pub fn new( + f: R, + timeout: Option, + 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,7 +392,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, None, false); r.read_until(&ReadUntil::NBytes(2)).expect("2 bytes"); assert_eq!( ("".to_string(), "rem ipsum dolor sit amet".to_string()), @@ -384,10 +400,21 @@ mod tests { ); } + #[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"); diff --git a/src/session.rs b/src/session.rs index 467a5b9f..6704bdfa 100644 --- a/src/session.rs +++ b/src/session.rs @@ -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, timeout_ms: Option, strip_ansi_escape_codes: bool) -> Self { Self { writer: LineWriter::new(writer), - reader: NBReader::new(reader, timeout_ms), + reader: NBReader::new(reader, timeout_ms, strip_ansi_escape_codes), } } @@ -138,7 +138,7 @@ impl StreamSession { /// /// # fn main() { /// # || -> Result<(), Error> { - /// let mut s = spawn("cat", Some(1000))?; + /// let mut s = spawn("cat", Some(1000), false)?; /// s.send_line("hello, polly!")?; /// s.exp_any(vec![ReadUntil::String("hello".into()), /// ReadUntil::EOF])?; @@ -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), false)?; +/// 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, timeout_ms: Option, strip_ansi_escape_codes: bool) -> 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, timeout_ms, strip_ansi_escape_codes); Ok(Self { process, stream }) } } @@ -200,7 +218,7 @@ fn tokenize_command(program: &str) -> Result, Error> { /// a problem the program just hangs instead of exiting with an /// error message indicating where it stopped. /// For automation 30'000 (30s, the default in pexpect) is a good value. -pub fn spawn(program: &str, timeout_ms: Option) -> Result { +pub fn spawn(program: &str, timeout_ms: Option, strip_ansi_escape_codes: bool) -> Result { if program.is_empty() { return Err(Error::EmptyProgramName); } @@ -209,11 +227,11 @@ pub fn spawn(program: &str, timeout_ms: Option) -> Result) -> Result { +pub fn spawn_command(command: Command, timeout_ms: Option, strip_ansi_escape_codes: bool) -> Result { #[cfg(feature = "which")] { let _ = which::which(command.get_program())?; @@ -221,7 +239,7 @@ pub fn spawn_command(command: Command, timeout_ms: Option) -> Result Result<(), Error> { - /// let mut p = spawn_bash(Some(1000))?; + /// let mut p = spawn_bash(Some(1000), false)?; /// p.execute("cat <(echo ready) -", "ready")?; /// p.send_line("hans")?; /// p.exp_string("hans")?; @@ -352,7 +370,7 @@ impl Drop for PtyReplSession { /// Also: if you start a program you should use `execute` and not `send_line`. /// /// For an example see the README -pub fn spawn_bash(timeout: Option) -> Result { +pub fn spawn_bash(timeout: Option, strip_ansi_escape_codes: bool) -> Result { // unfortunately working with a temporary tmpfile is the only // way to guarantee that we are "in step" with the prompt // all other attempts were futile, especially since we cannot @@ -372,7 +390,7 @@ pub fn spawn_bash(timeout: Option) -> Result { "--rcfile", rcfile.path().to_str().unwrap_or("temp file does not exist"), ]); - spawn_command(c, timeout).and_then(|p| { + spawn_command(c, timeout, strip_ansi_escape_codes).and_then(|p| { let new_prompt = "[REXPECT_PROMPT>"; let mut pb = PtyReplSession { prompt: new_prompt.to_string(), @@ -392,8 +410,8 @@ pub fn spawn_bash(timeout: Option) -> Result { /// Spawn the python shell /// /// This is just a proof of concept implementation (and serves for documentation purposes) -pub fn spawn_python(timeout: Option) -> Result { - spawn_command(Command::new("python"), timeout).map(|p| PtyReplSession { +pub fn spawn_python(timeout: Option, strip_ansi_escape_codes: bool) -> Result { + spawn_command(Command::new("python"), timeout, strip_ansi_escape_codes).map(|p| PtyReplSession { prompt: ">>> ".to_string(), pty_session: p, quit_command: Some("exit()".to_string()), @@ -407,7 +425,7 @@ pub fn spawn_stream( writer: W, timeout_ms: Option, ) -> StreamSession { - StreamSession::new(reader, writer, timeout_ms) + StreamSession::new(reader, writer, timeout_ms, false) } #[cfg(test)] @@ -416,7 +434,7 @@ mod tests { #[test] fn test_read_line() -> Result<(), Error> { - let mut s = spawn("cat", Some(100000))?; + let mut s = spawn("cat", Some(100000), false)?; s.send_line("hans")?; assert_eq!("hans", s.read_line()?); let should = crate::process::wait::WaitStatus::Signaled( @@ -430,7 +448,7 @@ mod tests { #[test] fn test_expect_eof_timeout() -> Result<(), Error> { - let mut p = spawn("sleep 3", Some(1000)).expect("cannot run sleep 3"); + let mut p = spawn("sleep 3", Some(1000), false).expect("cannot run sleep 3"); match p.exp_eof() { Ok(_) => panic!("should raise Timeout"), Err(Error::Timeout { .. }) => {} @@ -441,13 +459,13 @@ mod tests { #[test] fn test_expect_eof_timeout2() { - let mut p = spawn("sleep 1", Some(1100)).expect("cannot run sleep 1"); + let mut p = spawn("sleep 1", Some(1100), false).expect("cannot run sleep 1"); assert!(p.exp_eof().is_ok(), "expected eof"); } #[test] fn test_expect_string() -> Result<(), Error> { - let mut p = spawn("cat", Some(1000)).expect("cannot run cat"); + let mut p = spawn("cat", Some(1000), false).expect("cannot run cat"); p.send_line("hello world!")?; p.exp_string("hello world!")?; p.send_line("hello heaven!")?; @@ -457,7 +475,7 @@ mod tests { #[test] fn test_read_string_before() -> Result<(), Error> { - let mut p = spawn("cat", Some(1000)).expect("cannot run cat"); + let mut p = spawn("cat", Some(1000), false).expect("cannot run cat"); p.send_line("lorem ipsum dolor sit amet")?; assert_eq!("lorem ipsum dolor sit ", p.exp_string("amet")?); Ok(()) @@ -465,7 +483,7 @@ mod tests { #[test] fn test_expect_any() -> Result<(), Error> { - let mut p = spawn("cat", Some(1000)).expect("cannot run cat"); + let mut p = spawn("cat", Some(1000), false).expect("cannot run cat"); p.send_line("Hi")?; match p.exp_any(vec![ ReadUntil::NBytes(3), @@ -479,7 +497,7 @@ mod tests { #[test] fn test_expect_empty_command_error() { - let p = spawn("", Some(1000)); + let p = spawn("", Some(1000), false); match p { Ok(_) => panic!("should raise an error"), Err(Error::EmptyProgramName) => {} @@ -489,7 +507,7 @@ mod tests { #[test] fn test_kill_timeout() -> Result<(), Error> { - let mut p = spawn_bash(Some(1000))?; + let mut p = spawn_bash(Some(1000), false)?; p.execute("cat <(echo ready) -", "ready")?; Ok(()) // p is dropped here and kill is sent immediately to bash @@ -498,7 +516,7 @@ mod tests { #[test] fn test_bash() -> Result<(), Error> { - let mut p = spawn_bash(Some(1000))?; + let mut p = spawn_bash(Some(1000), false)?; p.send_line("cd /tmp/")?; p.wait_for_prompt()?; p.send_line("pwd")?; @@ -508,7 +526,7 @@ mod tests { #[test] fn test_bash_control_chars() -> Result<(), Error> { - let mut p = spawn_bash(Some(1000))?; + let mut p = spawn_bash(Some(1000), false)?; p.execute("cat <(echo ready) -", "ready")?; p.send_control('c')?; // abort: SIGINT p.wait_for_prompt()?; From 2771b71f2871fc9542d9316f2ebdb40e53972396 Mon Sep 17 00:00:00 2001 From: phaer Date: Sun, 23 Apr 2023 23:01:00 +0200 Subject: [PATCH 2/4] Add Options to configure NBReader Signed-off-by: phaer --- README.md | 6 ++-- examples/bash.rs | 2 +- examples/bash_read.rs | 2 +- examples/exit_code.rs | 4 +-- examples/ftp.rs | 2 +- examples/repl.rs | 2 +- src/lib.rs | 6 ++-- src/reader.rs | 68 ++++++++++++++++++++++----------------- src/session.rs | 74 +++++++++++++++++++++++++++---------------- 9 files changed, 97 insertions(+), 69 deletions(-) diff --git a/README.md b/README.md index 269b168e..0bac2320 100644 --- a/README.md +++ b/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), false)?; + let mut p = spawn("ftp speedtest.tele2.net", Some(30_000))?; 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), false)?; + let mut p = spawn_bash(Some(2000))?; // 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), false)?; + let mut p = spawn_bash(Some(1000))?; p.execute("ping 8.8.8.8", "bytes of data")?; p.send_control('z')?; p.wait_for_prompt()?; diff --git a/examples/bash.rs b/examples/bash.rs index 57e7677e..a95ac46b 100644 --- a/examples/bash.rs +++ b/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), false)?; + let mut p = spawn_bash(Some(1000))?; p.execute("ping 8.8.8.8", "bytes")?; p.send_control('z')?; p.wait_for_prompt()?; diff --git a/examples/bash_read.rs b/examples/bash_read.rs index fb3bf4c6..4522a062 100644 --- a/examples/bash_read.rs +++ b/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), false)?; + let mut p = spawn_bash(Some(2000))?; // case 1: wait until program is done p.send_line("hostname")?; diff --git a/examples/exit_code.rs b/examples/exit_code.rs index 9c52bacb..0964f6f8 100644 --- a/examples/exit_code.rs +++ b/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), false)?; + let p = spawn("cat /etc/passwd", Some(2000))?; 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), false)?; + let mut p = spawn("cat /this/does/not/exist", Some(2000))?; match p.process.wait() { Ok(wait::WaitStatus::Exited(_, 0)) => println!("cat succeeded"), Ok(wait::WaitStatus::Exited(_, c)) => { diff --git a/examples/ftp.rs b/examples/ftp.rs index 0ace5dd8..d1571879 100644 --- a/examples/ftp.rs +++ b/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), false)?; + let mut p = spawn("ftp speedtest.tele2.net", Some(2000))?; p.exp_regex("Name \\(.*\\):")?; p.send_line("anonymous")?; p.exp_string("Password")?; diff --git a/examples/repl.rs b/examples/repl.rs index 54e2e40d..5c0c257f 100644 --- a/examples/repl.rs +++ b/examples/repl.rs @@ -12,7 +12,7 @@ fn ed_session() -> Result { // used for `wait_for_prompt()` prompt: "> ".to_string(), - pty_session: spawn("/bin/ed -p '> '", Some(2000), false)?, + pty_session: spawn("/bin/ed -p '> '", Some(2000))?, // 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 diff --git a/src/lib.rs b/src/lib.rs index 9e9e1a82..96244abd 100644 --- a/src/lib.rs +++ b/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), false)?; +//! let mut p = spawn("ftp speedtest.tele2.net", Some(2000))?; //! 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), false)?; +//! let mut p = spawn_bash(Some(30_000))?; //! p.execute("ping 8.8.8.8", "bytes of data")?; //! p.send_control('z')?; //! p.wait_for_prompt()?; @@ -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 d013584c..83e05d55 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,15 +128,8 @@ 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 - /// - strip_ansi_escape_codes: Whether to filter out escape codes, such as colors. - pub fn new( - f: R, - timeout: Option, - strip_ansi_escape_codes: bool - ) -> 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 @@ -140,21 +145,20 @@ impl NBReader { break; } Ok(_) => { - if strip_ansi_escape_codes && byte[0] == 27 { + if options.strip_ansi_escape_codes && byte[0] == 27 { while let Ok(_) = reader.read(&mut byte) { if char::from(byte[0]).is_alphabetic() { break; } } - } - else { + } else { tx.send(Ok(PipedChar::Char(byte[0]))) - .map_err(|_| Error::MpscSendError)?; + .map_err(|_| Error::MpscSendError)?; } } Err(error) => { tx.send(Err(PipeError::IO(error))) - .map_err(|_| Error::MpscSendError)?; + .map_err(|_| Error::MpscSendError)?; } } } @@ -168,7 +172,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), } } @@ -219,11 +223,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, false); + /// 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); @@ -303,7 +307,7 @@ mod tests { #[test] fn test_expect_melon() { let f = io::Cursor::new("a melon\r\n"); - let mut r = NBReader::new(f, None, false); + 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())) @@ -320,7 +324,7 @@ mod tests { #[test] fn test_regex() { let f = io::Cursor::new("2014-03-15"); - let mut r = NBReader::new(f, None, false); + 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()), @@ -332,7 +336,7 @@ mod tests { #[test] fn test_regex2() { let f = io::Cursor::new("2014-03-15"); - let mut r = NBReader::new(f, None, false); + let mut r = NBReader::new(f, Options::default()); let re = Regex::new(r"-\d{2}-").unwrap(); assert_eq!( ("2014".to_string(), "-03-".to_string()), @@ -344,7 +348,7 @@ mod tests { #[test] fn test_nbytes() { let f = io::Cursor::new("abcdef"); - let mut r = NBReader::new(f, None, false); + let mut r = NBReader::new(f, Options::default()); assert_eq!( ("".to_string(), "ab".to_string()), r.read_until(&ReadUntil::NBytes(2)).expect("2 bytes") @@ -362,7 +366,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, false); + let mut r = NBReader::new(f, Options::default()); let result = r .read_until(&ReadUntil::Any(vec![ @@ -377,7 +381,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, false); + let mut r = NBReader::new(f, Options::default()); let result = r .read_until(&ReadUntil::Any(vec![ @@ -392,7 +396,7 @@ mod tests { #[test] fn test_eof() { let f = io::Cursor::new("lorem ipsum dolor sit amet"); - let mut r = NBReader::new(f, None, false); + 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()), @@ -403,18 +407,24 @@ mod tests { #[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(); + 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, false); + 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 6704bdfa..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, strip_ansi_escape_codes: bool) -> Self { + pub fn new(reader: R, writer: W, options: Options) -> Self { Self { writer: LineWriter::new(writer), - reader: NBReader::new(reader, timeout_ms, strip_ansi_escape_codes), + reader: NBReader::new(reader, options), } } @@ -138,7 +138,7 @@ impl StreamSession { /// /// # fn main() { /// # || -> Result<(), Error> { - /// let mut s = spawn("cat", Some(1000), false)?; + /// let mut s = spawn("cat", Some(1000))?; /// s.send_line("hello, polly!")?; /// s.exp_any(vec![ReadUntil::String("hello".into()), /// ReadUntil::EOF])?; @@ -182,7 +182,7 @@ impl DerefMut for PtySession { /// /// # fn main() { /// # || -> Result<(), Error> { -/// let mut s = spawn("cat", Some(1000), false)?; +/// let mut s = spawn("cat", Some(1000))?; /// s.send_line("hello, polly!")?; /// let line = s.read_line()?; /// assert_eq!("hello, polly!", line); @@ -191,10 +191,10 @@ impl DerefMut for PtySession { /// # } /// ``` impl PtySession { - fn new(process: PtyProcess, timeout_ms: Option, strip_ansi_escape_codes: bool) -> 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, strip_ansi_escape_codes); + let stream = StreamSession::new(reader, f, options); Ok(Self { process, stream }) } } @@ -218,7 +218,7 @@ fn tokenize_command(program: &str) -> Result, Error> { /// a problem the program just hangs instead of exiting with an /// error message indicating where it stopped. /// For automation 30'000 (30s, the default in pexpect) is a good value. -pub fn spawn(program: &str, timeout_ms: Option, strip_ansi_escape_codes: bool) -> Result { +pub fn spawn(program: &str, timeout_ms: Option) -> Result { if program.is_empty() { return Err(Error::EmptyProgramName); } @@ -227,19 +227,30 @@ pub fn spawn(program: &str, timeout_ms: Option, strip_ansi_escape_codes: bo let prog = parts.remove(0); let mut command = Command::new(prog); command.args(parts); - spawn_command(command, timeout_ms, strip_ansi_escape_codes) + spawn_command(command, timeout_ms) } /// See `spawn` -pub fn spawn_command(command: Command, timeout_ms: Option, strip_ansi_escape_codes: bool) -> Result { +pub fn spawn_command(command: Command, timeout_ms: Option) -> 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, strip_ansi_escape_codes) + PtySession::new(process, options) } /// A repl session: e.g. bash or the python shell: @@ -291,7 +302,7 @@ impl PtyReplSession { /// /// # fn main() { /// # || -> Result<(), Error> { - /// let mut p = spawn_bash(Some(1000), false)?; + /// let mut p = spawn_bash(Some(1000))?; /// p.execute("cat <(echo ready) -", "ready")?; /// p.send_line("hans")?; /// p.exp_string("hans")?; @@ -370,7 +381,7 @@ impl Drop for PtyReplSession { /// Also: if you start a program you should use `execute` and not `send_line`. /// /// For an example see the README -pub fn spawn_bash(timeout: Option, strip_ansi_escape_codes: bool) -> Result { +pub fn spawn_bash(timeout: Option) -> Result { // unfortunately working with a temporary tmpfile is the only // way to guarantee that we are "in step" with the prompt // all other attempts were futile, especially since we cannot @@ -390,7 +401,7 @@ pub fn spawn_bash(timeout: Option, strip_ansi_escape_codes: bool) -> Result "--rcfile", rcfile.path().to_str().unwrap_or("temp file does not exist"), ]); - spawn_command(c, timeout, strip_ansi_escape_codes).and_then(|p| { + spawn_command(c, timeout).and_then(|p| { let new_prompt = "[REXPECT_PROMPT>"; let mut pb = PtyReplSession { prompt: new_prompt.to_string(), @@ -410,8 +421,8 @@ pub fn spawn_bash(timeout: Option, strip_ansi_escape_codes: bool) -> Result /// Spawn the python shell /// /// This is just a proof of concept implementation (and serves for documentation purposes) -pub fn spawn_python(timeout: Option, strip_ansi_escape_codes: bool) -> Result { - spawn_command(Command::new("python"), timeout, strip_ansi_escape_codes).map(|p| PtyReplSession { +pub fn spawn_python(timeout: Option) -> Result { + spawn_command(Command::new("python"), timeout).map(|p| PtyReplSession { prompt: ">>> ".to_string(), pty_session: p, quit_command: Some("exit()".to_string()), @@ -425,7 +436,14 @@ pub fn spawn_stream( writer: W, timeout_ms: Option, ) -> StreamSession { - StreamSession::new(reader, writer, timeout_ms, false) + StreamSession::new( + reader, + writer, + Options { + timeout_ms, + strip_ansi_escape_codes: false, + }, + ) } #[cfg(test)] @@ -434,7 +452,7 @@ mod tests { #[test] fn test_read_line() -> Result<(), Error> { - let mut s = spawn("cat", Some(100000), false)?; + let mut s = spawn("cat", Some(100000))?; s.send_line("hans")?; assert_eq!("hans", s.read_line()?); let should = crate::process::wait::WaitStatus::Signaled( @@ -448,7 +466,7 @@ mod tests { #[test] fn test_expect_eof_timeout() -> Result<(), Error> { - let mut p = spawn("sleep 3", Some(1000), false).expect("cannot run sleep 3"); + let mut p = spawn("sleep 3", Some(1000)).expect("cannot run sleep 3"); match p.exp_eof() { Ok(_) => panic!("should raise Timeout"), Err(Error::Timeout { .. }) => {} @@ -459,13 +477,13 @@ mod tests { #[test] fn test_expect_eof_timeout2() { - let mut p = spawn("sleep 1", Some(1100), false).expect("cannot run sleep 1"); + let mut p = spawn("sleep 1", Some(1100)).expect("cannot run sleep 1"); assert!(p.exp_eof().is_ok(), "expected eof"); } #[test] fn test_expect_string() -> Result<(), Error> { - let mut p = spawn("cat", Some(1000), false).expect("cannot run cat"); + let mut p = spawn("cat", Some(1000)).expect("cannot run cat"); p.send_line("hello world!")?; p.exp_string("hello world!")?; p.send_line("hello heaven!")?; @@ -475,7 +493,7 @@ mod tests { #[test] fn test_read_string_before() -> Result<(), Error> { - let mut p = spawn("cat", Some(1000), false).expect("cannot run cat"); + let mut p = spawn("cat", Some(1000)).expect("cannot run cat"); p.send_line("lorem ipsum dolor sit amet")?; assert_eq!("lorem ipsum dolor sit ", p.exp_string("amet")?); Ok(()) @@ -483,7 +501,7 @@ mod tests { #[test] fn test_expect_any() -> Result<(), Error> { - let mut p = spawn("cat", Some(1000), false).expect("cannot run cat"); + let mut p = spawn("cat", Some(1000)).expect("cannot run cat"); p.send_line("Hi")?; match p.exp_any(vec![ ReadUntil::NBytes(3), @@ -497,7 +515,7 @@ mod tests { #[test] fn test_expect_empty_command_error() { - let p = spawn("", Some(1000), false); + let p = spawn("", Some(1000)); match p { Ok(_) => panic!("should raise an error"), Err(Error::EmptyProgramName) => {} @@ -507,7 +525,7 @@ mod tests { #[test] fn test_kill_timeout() -> Result<(), Error> { - let mut p = spawn_bash(Some(1000), false)?; + let mut p = spawn_bash(Some(1000))?; p.execute("cat <(echo ready) -", "ready")?; Ok(()) // p is dropped here and kill is sent immediately to bash @@ -516,7 +534,7 @@ mod tests { #[test] fn test_bash() -> Result<(), Error> { - let mut p = spawn_bash(Some(1000), false)?; + let mut p = spawn_bash(Some(1000))?; p.send_line("cd /tmp/")?; p.wait_for_prompt()?; p.send_line("pwd")?; @@ -526,7 +544,7 @@ mod tests { #[test] fn test_bash_control_chars() -> Result<(), Error> { - let mut p = spawn_bash(Some(1000), false)?; + let mut p = spawn_bash(Some(1000))?; p.execute("cat <(echo ready) -", "ready")?; p.send_control('c')?; // abort: SIGINT p.wait_for_prompt()?; From 65d656a546ede3673940cbff0a2ad88294392069 Mon Sep 17 00:00:00 2001 From: phaer Date: Wed, 26 Apr 2023 16:43:43 +0200 Subject: [PATCH 3/4] Add test case for partial ansi codes Signed-off-by: phaer --- src/reader.rs | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/reader.rs b/src/reader.rs index 83e05d55..5215dd0d 100644 --- a/src/reader.rs +++ b/src/reader.rs @@ -404,6 +404,23 @@ 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"); From e992f993865ba56a1ba53092f920802b37f56b95 Mon Sep 17 00:00:00 2001 From: phaer Date: Thu, 27 Apr 2023 12:49:24 +0200 Subject: [PATCH 4/4] Remove nested loop to strip ansi codes because handling of Err(_) and Ok(0) might be tricky otherwise, so we just use a boolean flag Signed-off-by: phaer --- src/reader.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/reader.rs b/src/reader.rs index 5215dd0d..1f5d8b62 100644 --- a/src/reader.rs +++ b/src/reader.rs @@ -136,6 +136,7 @@ impl NBReader { 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) { @@ -146,10 +147,10 @@ impl NBReader { } Ok(_) => { if options.strip_ansi_escape_codes && byte[0] == 27 { - while let Ok(_) = reader.read(&mut byte) { - if char::from(byte[0]).is_alphabetic() { - break; - } + 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])))