From 0212cd6ff069e5633ea8e6d4dad8212f88ef6464 Mon Sep 17 00:00:00 2001 From: Maxim Zhiburt Date: Tue, 26 May 2020 22:32:28 +0300 Subject: [PATCH 1/4] draft of a new API Signed-off-by: Maxim Zhiburt --- src/reader.rs | 268 ++++++++++++++++++++++++++++++++++++++----------- src/session.rs | 115 ++++++++++++--------- 2 files changed, 279 insertions(+), 104 deletions(-) diff --git a/src/reader.rs b/src/reader.rs index f479991b..da4d4462 100644 --- a/src/reader.rs +++ b/src/reader.rs @@ -1,12 +1,12 @@ //! Unblocking reader which supports waiting for strings/regexes and EOF to be present -use std::io::{self, BufReader}; -use std::io::prelude::*; -use std::sync::mpsc::{channel, Receiver}; -use std::{thread, result}; -use std::{time, fmt}; use crate::errors::*; // load error-chain pub use regex::Regex; +use std::io::prelude::*; +use std::io::{self, BufReader}; +use std::sync::mpsc::{channel, Receiver}; +use std::{fmt, time}; +use std::{result, thread}; #[derive(Debug)] enum PipeError { @@ -19,12 +19,42 @@ enum PipedChar { EOF, } +#[derive(Clone)] pub enum ReadUntil { String(String), Regex(Regex), EOF, NBytes(usize), - Any(Vec), +} + +impl ReadUntil { + pub fn string_needle(self) -> Option> { + match self { + ReadUntil::String(s) => Some(s), + _ => None + } + } + + pub fn eof_needle(self) -> Option> { + match self { + ReadUntil::EOF => Some(EOF), + _ => None + } + } + + pub fn nbytes_needle(self) -> Option> { + match self { + ReadUntil::NBytes(n) => Some(NBytes(n)), + _ => None + } + } + + pub fn regex_needle(self) -> Option> { + match self { + ReadUntil::Regex(regex) => Some(regex), + _ => None + } + } } impl fmt::Display for ReadUntil { @@ -36,13 +66,6 @@ impl fmt::Display for ReadUntil { &ReadUntil::Regex(ref r) => format!("Regex: \"{}\"", r), &ReadUntil::EOF => "EOF (End of File)".to_string(), &ReadUntil::NBytes(n) => format!("reading {} bytes", n), - &ReadUntil::Any(ref v) => { - let mut res = Vec::new(); - for r in v { - res.push(r.to_string()); - } - res.join(", ") - } }; write!(f, "{}", printable) } @@ -70,7 +93,13 @@ pub fn find(needle: &ReadUntil, buffer: &str, eof: bool) -> Option<(usize, usize None } } - &ReadUntil::EOF => if eof { Some((0, buffer.len())) } else { None }, + &ReadUntil::EOF => { + if eof { + Some((0, buffer.len())) + } else { + None + } + } &ReadUntil::NBytes(n) => { if n <= buffer.len() { Some((0, n)) @@ -82,17 +111,127 @@ pub fn find(needle: &ReadUntil, buffer: &str, eof: bool) -> Option<(usize, usize None } } - &ReadUntil::Any(ref any) => { - for read_until in any { - if let Some(pos_tuple) = find(&read_until, buffer, eof) { - return Some(pos_tuple); - } - } + } +} + +pub struct Match { + begin: usize, + end: usize, + pub interest: T, +} + +impl Match { + fn new(begin: usize, end: usize, obj: T) -> Self { + Self { + begin, + end, + interest: obj, + } + } +} + +pub trait Needle { + type Interest; + + fn find(&self, buffer: &str, eof: bool) -> Option>; +} + +impl Needle for str { + type Interest = String; + + fn find(&self, buffer: &str, eof: bool) -> Option> { + buffer.find(self).and_then(|pos| { + Some(Match::new( + pos, + pos + self.len(), + buffer[..pos].to_string(), + )) + }) + } +} + +impl Needle for String { + type Interest = String; + + fn find(&self, buffer: &str, eof: bool) -> Option> { + buffer.find(self).and_then(|pos| { + Some(Match::new( + pos, + pos + self.len(), + buffer[..pos].to_string(), + )) + }) + } +} + +impl Needle for Regex { + type Interest = (String, String); + + fn find(&self, buffer: &str, eof: bool) -> Option> { + if let Some(mat) = Regex::find(&self, buffer) { + Some(Match::new( + 0, + mat.end(), + ( + buffer[..mat.start()].to_string(), + buffer[mat.start()..mat.end()].to_string(), + ), + )) + } else { None } } } +pub struct EOF; + +impl Needle for EOF { + type Interest = String; + + fn find(&self, buffer: &str, eof: bool) -> Option> { + if eof { + Some(Match::new(0, buffer.len(), buffer.to_string())) + } else { + None + } + } +} + +pub struct NBytes(pub usize); + +impl Needle for NBytes { + type Interest = String; + + fn find(&self, buffer: &str, eof: bool) -> Option> { + if self.0 <= buffer.len() { + Some(Match::new(0, self.0, buffer[..self.0].to_string())) + } else if eof && buffer.len() > 0 { + // reached almost end of buffer, return string, even though it will be + // smaller than the wished n bytes + Some(Match::new(0, buffer.len(), buffer[..buffer.len()].to_string())) + } else { + None + } + } +} + +impl Needle for [ReadUntil] { + type Interest = (String, usize); + + fn find(&self, buffer: &str, eof: bool) -> Option> { + self.iter() + .enumerate() + .find_map(|(index, needle)| match find(needle, buffer, eof) { + Some((start, end)) => Some(Match::new( + start, + end, + (buffer.to_string(), index), + )), + None => None, + }) + } +} + /// Non blocking reader /// /// Typically you'd need that to check for output of a process without blocking your thread. @@ -200,55 +339,60 @@ impl NBReader { /// /// ``` /// # use std::io::Cursor; - /// use rexpect::reader::{NBReader, ReadUntil, Regex}; + /// use rexpect::reader::{NBReader, ReadUntil, NBytes, EOF, Regex}; /// // 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 (first_line, _) = e.read_until(&ReadUntil::String('\n'.to_string())).unwrap(); + /// let first_line = e.read_until("\n").unwrap(); /// assert_eq!("Hello, miss!", &first_line); /// - /// let (_, two_bytes) = e.read_until(&ReadUntil::NBytes(2)).unwrap(); + /// let two_bytes = e.read_until(&NBytes(2)).unwrap(); /// assert_eq!("Wh", &two_bytes); /// /// let re = Regex::new(r"'[a-z]+'").unwrap(); // will find 'miss' - /// let (before, reg_match) = e.read_until(&ReadUntil::Regex(re)).unwrap(); + /// let (before, reg_match) = e.read_until(&re).unwrap(); /// assert_eq!("at do you mean: ", &before); /// assert_eq!("'miss'", ®_match); /// - /// let (_, until_end) = e.read_until(&ReadUntil::EOF).unwrap(); + /// let until_end = e.read_until(&EOF).unwrap(); /// assert_eq!("?", &until_end); /// ``` /// - pub fn read_until(&mut self, needle: &ReadUntil) -> Result<(String, String)> { + pub fn read_until(&mut self, needle: &N) -> Result { let start = time::Instant::now(); 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(); - return Ok((first, second)); + if let Some(m) = needle.find(&self.buffer, self.eof) { + self.buffer.drain(..m.begin); + self.buffer.drain(..m.end - m.begin); + return Ok(m.interest); } // reached end of stream and didn't match -> error // we don't know the reason of eof yet, so we provide an empty string // this will be filled out in session::exp() if self.eof { - return Err(ErrorKind::EOF(needle.to_string(), self.buffer.clone(), None).into()); + return Err( + ErrorKind::EOF("ERROR NEEDLE".to_string(), self.buffer.clone(), None).into(), + ); } // ran into timeout if let Some(timeout) = self.timeout { if start.elapsed() > timeout { - return Err(ErrorKind::Timeout(needle.to_string(), - self.buffer.clone() - .replace("\n", "`\\n`\n") - .replace("\r", "`\\r`") - .replace('\u{1b}', "`^`"), - timeout) - .into()); + return Err(ErrorKind::Timeout( + "ERROR NEEDLE".to_string(), + self.buffer + .clone() + .replace("\n", "`\\n`\n") + .replace("\r", "`\\r`") + .replace('\u{1b}', "`^`"), + timeout, + ) + .into()); } } // nothing matched: wait a little @@ -277,11 +421,12 @@ mod tests { fn test_expect_melon() { let f = io::Cursor::new("a melon\r\n"); let mut r = NBReader::new(f, None); - assert_eq!(("a melon".to_string(), "\r\n".to_string()), - r.read_until(&ReadUntil::String("\r\n".to_string())) - .expect("cannot read line")); + assert_eq!( + "a melon".to_string(), + r.read_until("\r\n").expect("cannot read line") + ); // check for EOF - match r.read_until(&ReadUntil::NBytes(10)) { + match r.read_until(&NBytes(10)) { Ok(_) => assert!(false), Err(Error(ErrorKind::EOF(_, _, _), _)) => {} Err(Error(_, _)) => assert!(false), @@ -293,8 +438,7 @@ mod tests { let f = io::Cursor::new("2014-03-15"); let mut r = NBReader::new(f, None); let re = Regex::new(r"^\d{4}-\d{2}-\d{2}$").unwrap(); - r.read_until(&ReadUntil::Regex(re)) - .expect("regex doesn't match"); + r.read_until(&re).expect("regex doesn't match"); } #[test] @@ -302,42 +446,50 @@ mod tests { let f = io::Cursor::new("2014-03-15"); let mut r = NBReader::new(f, None); let re = Regex::new(r"-\d{2}-").unwrap(); - assert_eq!(("2014".to_string(), "-03-".to_string()), - r.read_until(&ReadUntil::Regex(re)) - .expect("regex doesn't match")); + assert_eq!( + ("2014".to_string(), "-03-".to_string()), + r.read_until(&re).expect("regex doesn't match") + ); } #[test] fn test_nbytes() { let f = io::Cursor::new("abcdef"); let mut r = NBReader::new(f, None); - assert_eq!(("".to_string(), "ab".to_string()), - r.read_until(&ReadUntil::NBytes(2)).expect("2 bytes")); - assert_eq!(("".to_string(), "cde".to_string()), - r.read_until(&ReadUntil::NBytes(3)).expect("3 bytes")); - assert_eq!(("".to_string(), "f".to_string()), - r.read_until(&ReadUntil::NBytes(4)).expect("4 bytes")); + assert_eq!( + "ab".to_string(), + r.read_until(&NBytes(2)).expect("2 bytes") + ); + assert_eq!( + "cde".to_string(), + r.read_until(&NBytes(3)).expect("3 bytes") + ); + assert_eq!( + "f".to_string(), + r.read_until(&NBytes(4)).expect("4 bytes") + ); } #[test] fn test_eof() { let f = io::Cursor::new("lorem ipsum dolor sit amet"); let mut r = NBReader::new(f, None); - 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")); + r.read_until(&NBytes(2)).expect("2 bytes"); + assert_eq!( + "rem ipsum dolor sit amet".to_string(), + r.read_until(&EOF).expect("reading until EOF") + ); } #[test] fn test_try_read() { let f = io::Cursor::new("lorem"); let mut r = NBReader::new(f, None); - r.read_until(&ReadUntil::NBytes(4)).expect("4 bytes"); + r.read_until(&NBytes(4)).expect("4 bytes"); assert_eq!(Some('m'), r.try_read()); assert_eq!(None, r.try_read()); assert_eq!(None, r.try_read()); assert_eq!(None, r.try_read()); assert_eq!(None, r.try_read()); } - } diff --git a/src/session.rs b/src/session.rs index e01b3a3c..8c98422e 100644 --- a/src/session.rs +++ b/src/session.rs @@ -1,14 +1,14 @@ //! Main module of rexpect: start new process and interact with it +use crate::errors::*; // load error-chain use crate::process::PtyProcess; -use crate::reader::{NBReader, Regex}; pub use crate::reader::ReadUntil; +use crate::reader::{self, NBReader, Needle, Regex}; use std::fs::File; -use std::io::LineWriter; -use std::process::Command; use std::io::prelude::*; +use std::io::LineWriter; use std::ops::{Deref, DerefMut}; -use crate::errors::*; // load error-chain +use std::process::Command; use tempfile; /// Interact with a process with read/write/signals, etc. @@ -46,13 +46,13 @@ impl PtySession { /// returns number of written bytes pub fn send_line(&mut self, line: &str) -> Result { let mut len = self.send(line)?; - len += self.writer + len += self + .writer .write(&['\n' as u8]) .chain_err(|| "cannot write newline")?; Ok(len) } - /// Send string to process. As stdin of the process is most likely buffered, you'd /// need to call `flush()` after `send()` to make the process actually see your input. /// @@ -89,7 +89,7 @@ impl PtySession { } // wrapper around reader::read_until to give more context for errors - fn exp(&mut self, needle: &ReadUntil) -> Result<(String, String)> { + pub fn exp(&mut self, needle: &N) -> Result { match self.reader.read_until(needle) { Ok(s) => Ok(s), Err(Error(ErrorKind::EOF(expected, got, _), _)) => { @@ -107,8 +107,8 @@ impl PtySession { /// Read one line (blocking!) and return line without the newline /// (waits until \n is in the output fetches the line and removes \r at the end if present) pub fn read_line(&mut self) -> Result { - match self.exp(&ReadUntil::String('\n'.to_string())) { - Ok((mut line, _)) => { + match self.exp("\n") { + Ok(mut line) => { if line.ends_with('\r') { line.pop().expect("this never happens"); } @@ -127,7 +127,7 @@ impl PtySession { /// Wait until we see EOF (i.e. child process has terminated) /// Return all the yet unread output pub fn exp_eof(&mut self) -> Result { - self.exp(&ReadUntil::EOF).and_then(|(_, s)| Ok(s)) + self.exp(&reader::EOF).and_then(|s| Ok(s)) } /// Wait until provided regex is seen on stdout of child process. @@ -138,23 +138,20 @@ impl PtySession { /// Note that `exp_regex("^foo")` matches the start of the yet consumed output. /// For matching the start of the line use `exp_regex("\nfoo")` pub fn exp_regex(&mut self, regex: &str) -> Result<(String, String)> { - let res = self.exp(&ReadUntil::Regex(Regex::new(regex).chain_err(|| "invalid regex")?)) - .and_then(|s| Ok(s)); - res + let res = self.exp(&Regex::new(regex).chain_err(|| "invalid regex")?); + res.map(|res| (res.0, res.1)) } /// Wait until provided string is seen on stdout of child process. /// Return the yet unread output (without the matched string) pub fn exp_string(&mut self, needle: &str) -> Result { - self.exp(&ReadUntil::String(needle.to_string())) - .and_then(|(s, _)| Ok(s)) + self.exp(needle).and_then(|s| Ok(s)) } /// Wait until provided char is seen on stdout of child process. /// Return the yet unread output (without the matched char) pub fn exp_char(&mut self, needle: char) -> Result { - self.exp(&ReadUntil::String(needle.to_string())) - .and_then(|(s, _)| Ok(s)) + self.exp(needle.to_string().as_str()).and_then(|s| Ok(s)) } /// Wait until any of the provided needles is found. @@ -179,8 +176,8 @@ impl PtySession { /// # }().expect("test failed"); /// # } /// ``` - pub fn exp_any(&mut self, needles: Vec) -> Result<(String, String)> { - self.exp(&ReadUntil::Any(needles)) + pub fn exp_any(&mut self, needles: Vec) -> Result<(String, usize)> { + self.exp(needles.as_slice()) } } @@ -224,19 +221,18 @@ pub fn spawn(program: &str, timeout_ms: Option) -> Result { /// See `spawn` pub fn spawn_command(command: Command, timeout_ms: Option) -> Result { let commandname = format!("{:?}", &command); - let mut process = PtyProcess::new(command) - .chain_err(|| "couldn't start process")?; + let mut process = PtyProcess::new(command).chain_err(|| "couldn't start process")?; process.set_kill_timeout(timeout_ms); let f = process.get_file_handle(); let writer = LineWriter::new(f.try_clone().chain_err(|| "couldn't open write stream")?); let reader = NBReader::new(f, timeout_ms); Ok(PtySession { - process: process, - writer: writer, - reader: reader, - commandname: commandname, - }) + process: process, + writer: writer, + reader: reader, + commandname: commandname, + }) } /// A repl session: e.g. bash or the python shell: @@ -344,7 +340,6 @@ impl Drop for PtyReplSession { } } - /// Spawn bash in a pty session, run programs and expect output /// /// @@ -376,13 +371,23 @@ pub fn spawn_bash(timeout: Option) -> Result { // would set as PS1 and we cannot know when is the right time // to set the new PS1 let mut rcfile = tempfile::NamedTempFile::new().unwrap(); - rcfile.write(b"include () { [[ -f \"$1\" ]] && source \"$1\"; }\n\ + rcfile + .write( + b"include () { [[ -f \"$1\" ]] && source \"$1\"; }\n\ include /etc/bash.bashrc\n\ include ~/.bashrc\n\ PS1=\"~~~~\"\n\ - unset PROMPT_COMMAND\n").expect("cannot write to tmpfile"); + unset PROMPT_COMMAND\n", + ) + .expect("cannot write to tmpfile"); let mut c = Command::new("bash"); - c.args(&["--rcfile", rcfile.path().to_str().unwrap_or_else(|| return "temp file does not exist".into())]); + c.args(&[ + "--rcfile", + rcfile + .path() + .to_str() + .unwrap_or_else(|| return "temp file does not exist".into()), + ]); spawn_command(c, timeout).and_then(|p| { let new_prompt = "[REXPECT_PROMPT>"; let mut pb = PtyReplSession { @@ -392,7 +397,9 @@ pub fn spawn_bash(timeout: Option) -> Result { echo_on: false, }; pb.exp_string("~~~~")?; - rcfile.close().chain_err(|| "cannot delete temporary rcfile")?; + rcfile + .close() + .chain_err(|| "cannot delete temporary rcfile")?; pb.send_line(&("PS1='".to_string() + new_prompt + "'"))?; // wait until the new prompt appears pb.wait_for_prompt()?; @@ -424,16 +431,17 @@ mod tests { let mut s = spawn("cat", Some(1000))?; s.send_line("hans")?; assert_eq!("hans", s.read_line()?); - let should = crate::process::wait::WaitStatus::Signaled(s.process.child_pid, - crate::process::signal::Signal::SIGTERM, - false); + let should = crate::process::wait::WaitStatus::Signaled( + s.process.child_pid, + crate::process::signal::Signal::SIGTERM, + false, + ); assert_eq!(should, s.process.exit()?); Ok(()) }() - .unwrap_or_else(|e| panic!("test_read_line failed: {}", e)); + .unwrap_or_else(|e| panic!("test_read_line failed: {}", e)); } - #[test] fn test_expect_eof_timeout() { || -> Result<()> { @@ -446,7 +454,7 @@ mod tests { } Ok(()) }() - .unwrap_or_else(|e| panic!("test_timeout failed: {}", e)); + .unwrap_or_else(|e| panic!("test_timeout failed: {}", e)); } #[test] @@ -465,7 +473,7 @@ mod tests { p.exp_string("hello heaven!")?; Ok(()) }() - .unwrap_or_else(|e| panic!("test_expect_string failed: {}", e)); + .unwrap_or_else(|e| panic!("test_expect_string failed: {}", e)); } #[test] @@ -476,21 +484,35 @@ mod tests { assert_eq!("lorem ipsum dolor sit ", p.exp_string("amet")?); Ok(()) }() - .unwrap_or_else(|e| panic!("test_read_string_before failed: {}", e)); + .unwrap_or_else(|e| panic!("test_read_string_before failed: {}", e)); } #[test] fn test_expect_any() { || -> Result<()> { let mut p = spawn("cat", Some(1000)).expect("cannot run cat"); - p.send_line("Hi")?; - match p.exp_any(vec![ReadUntil::NBytes(3), ReadUntil::String("Hi".to_string())]) { - Ok(s) => assert_eq!(("".to_string(), "Hi\r".to_string()), s), + p.send_line("world Hi")?; + let until = vec![ReadUntil::NBytes(30), ReadUntil::String("Hi".to_string())]; + + match p.exp(until.as_slice()) { + Ok((buffer, index)) => { + // why the buffer contains doubled lines + // does it normal? + assert_eq!("world Hi\r\nworld Hi\r\n".to_string(), buffer); + assert_eq!(index, 1); + let res = &until[1] + .clone() + .string_needle() + .unwrap() + .find(&buffer, true) + .expect("unexpected needle"); + assert_eq!("world ".to_string(), res.interest); + } Err(e) => assert!(false, format!("got error: {}", e)), } Ok(()) }() - .unwrap_or_else(|e| panic!("test_expect_any failed: {}", e)); + .unwrap_or_else(|e| panic!("test_expect_any failed: {}", e)); } #[test] @@ -499,7 +521,8 @@ mod tests { let mut p = spawn_bash(Some(1000))?; p.execute("cat <(echo ready) -", "ready")?; Ok(()) - }().unwrap_or_else(|e| panic!("test_kill_timeout failed: {}", e)); + }() + .unwrap_or_else(|e| panic!("test_kill_timeout failed: {}", e)); // p is dropped here and kill is sent immediatly to bash // Since that is not enough to make bash exit, a kill -9 is sent within 1s (timeout) } @@ -514,7 +537,7 @@ mod tests { assert_eq!("/tmp\r\n", p.wait_for_prompt()?); Ok(()) }() - .unwrap_or_else(|e| panic!("test_bash failed: {}", e)); + .unwrap_or_else(|e| panic!("test_bash failed: {}", e)); } #[test] @@ -532,7 +555,7 @@ mod tests { p.send_control('c')?; Ok(()) }() - .unwrap_or_else(|e| panic!("test_bash_control_chars failed: {}", e)); + .unwrap_or_else(|e| panic!("test_bash_control_chars failed: {}", e)); } #[test] From f89abfcbb1b067cb7dc8d7b5f403b79059c25275 Mon Sep 17 00:00:00 2001 From: Philipp Keller <11523+philippkeller@users.noreply.github.com> Date: Sun, 31 May 2020 14:33:56 +0200 Subject: [PATCH 2/4] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 7afcdc38..d4f6e1a1 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ Add this to your `Cargo.toml` ```toml [dependencies] -rexpect = "0.3" +rexpect = "0.4" ``` Simple example for interacting via ftp: From b8dc264ddb9783c9fe8d6dbf6b868b1b2814e95e Mon Sep 17 00:00:00 2001 From: Maxim Zhiburt Date: Mon, 8 Jun 2020 00:13:14 +0300 Subject: [PATCH 3/4] Implementation v2 Signed-off-by: Maxim Zhiburt --- src/lib.rs | 1 - src/reader.rs | 263 ++++++++++++++++++++++++++++++++----------------- src/session.rs | 73 +++++++------- 3 files changed, 209 insertions(+), 128 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 23a6c731..79fc0271 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -83,7 +83,6 @@ pub mod session; pub mod reader; pub use session::{spawn, spawn_bash, spawn_python}; -pub use reader::ReadUntil; pub mod errors { use std::time; diff --git a/src/reader.rs b/src/reader.rs index f479991b..0329a8a8 100644 --- a/src/reader.rs +++ b/src/reader.rs @@ -19,80 +19,171 @@ enum PipedChar { EOF, } -pub enum ReadUntil { - String(String), - Regex(Regex), - EOF, - NBytes(usize), - Any(Vec), +pub struct Match { + begin: usize, + end: usize, + interest: T, } -impl fmt::Display for ReadUntil { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - let printable = match self { - &ReadUntil::String(ref s) if s == "\n" => "\\n (newline)".to_string(), - &ReadUntil::String(ref s) if s == "\r" => "\\r (carriage return)".to_string(), - &ReadUntil::String(ref s) => format!("\"{}\"", s), - &ReadUntil::Regex(ref r) => format!("Regex: \"{}\"", r), - &ReadUntil::EOF => "EOF (End of File)".to_string(), - &ReadUntil::NBytes(n) => format!("reading {} bytes", n), - &ReadUntil::Any(ref v) => { - let mut res = Vec::new(); - for r in v { - res.push(r.to_string()); - } - res.join(", ") - } +impl Match { + fn new(begin: usize, end: usize, obj: T) -> Self { + Self { + begin, + end, + interest: obj, + } + } +} + +pub trait Needle { + type Interest; + + fn find(&self, buffer: &str, eof: bool) -> Option>; +} + +pub struct Str>(pub S); + +impl> Needle for Str { + type Interest = String; + + fn find(&self, buffer: &str, eof: bool) -> Option> { + let s = self.0.as_ref(); + buffer.find(s).and_then(|pos| { + Some(Match::new( + pos, + pos + s.len(), + buffer[..pos].to_string(), + )) + }) + } +} + +impl> fmt::Display for Str { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let printable = match self.0.as_ref() { + "\n" => "\\n (newline)".to_string(), + "\r" => "\\r (carriage return)".to_string(), + _ => format!("\"{}\"", self), }; + write!(f, "{}", printable) } } -/// find first occurrence of needle within buffer -/// -/// # Arguments: -/// -/// - buffer: the currently read buffer from a process which will still grow in the future -/// - eof: if the process already sent an EOF or a HUP -/// -/// # Return -/// -/// Tuple with match positions: -/// 1. position before match (0 in case of EOF and Nbytes) -/// 2. position after match -pub fn find(needle: &ReadUntil, buffer: &str, eof: bool) -> Option<(usize, usize)> { - match needle { - &ReadUntil::String(ref s) => buffer.find(s).and_then(|pos| Some((pos, pos + s.len()))), - &ReadUntil::Regex(ref pattern) => { - if let Some(mat) = pattern.find(buffer) { - Some((mat.start(), mat.end())) - } else { - None - } +pub struct Regx(pub Regex); + +impl Needle for Regx { + type Interest = (String, String); + + fn find(&self, buffer: &str, eof: bool) -> Option> { + if let Some(mat) = Regex::find(&self.0, buffer) { + Some(Match::new( + 0, + mat.end(), + ( + buffer[..mat.start()].to_string(), + buffer[mat.start()..mat.end()].to_string(), + ), + )) + } else { + None } - &ReadUntil::EOF => if eof { Some((0, buffer.len())) } else { None }, - &ReadUntil::NBytes(n) => { - if n <= buffer.len() { - Some((0, n)) - } else if eof && buffer.len() > 0 { - // reached almost end of buffer, return string, even though it will be - // smaller than the wished n bytes - Some((0, buffer.len())) - } else { - None - } + } +} + +impl fmt::Display for Regx { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "Regex: \"{}\"", self.0) + } +} + +pub struct EOF; + +impl Needle for EOF { + type Interest = String; + + fn find(&self, buffer: &str, eof: bool) -> Option> { + if eof { + Some(Match::new(0, buffer.len(), buffer.to_string())) + } else { + None } - &ReadUntil::Any(ref any) => { - for read_until in any { - if let Some(pos_tuple) = find(&read_until, buffer, eof) { - return Some(pos_tuple); - } - } + } +} + +impl fmt::Display for EOF { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "EOF (End of File)") + } +} + +pub struct NBytes(pub usize); + +impl Needle for NBytes { + type Interest = String; + + fn find(&self, buffer: &str, eof: bool) -> Option> { + if self.0 <= buffer.len() { + Some(Match::new(0, self.0, buffer[..self.0].to_string())) + } else if eof && buffer.len() > 0 { + // reached almost end of buffer, return string, even though it will be + // smaller than the wished n bytes + Some(Match::new(0, buffer.len(), buffer[..buffer.len()].to_string())) + } else { None } } } +impl fmt::Display for NBytes { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "reading {} bytes", self.0) + } +} + +pub struct UntilOr(A, B); + +pub enum OrInterest { + Lhs(A), + Rhs(B), +} + +impl Needle for UntilOr { + type Interest = OrInterest; + + fn find(&self, buffer: &str, eof: bool) -> Option> { + self.0 + .find(buffer, eof) + .and_then( |m| Some(Match::new(m.begin, m.end, OrInterest::Lhs(m.interest)))) + .or_else(|| + self.1.find(buffer, eof) + .and_then(|m| Some(Match::new(m.begin, m.end, OrInterest::Rhs(m.interest)))) + ) + } +} + +impl fmt::Display for UntilOr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "until {} or {};", self.0, self.1) + } +} + +// unfortunately we can't implement Needle for such type +// so there's a AsRef implementation instead of +pub struct Until(pub N); + +impl Until { + pub fn or(self, lhs: R) -> Until> { + Until(UntilOr(self.0, lhs)) + } +} + +impl AsRef for Until { + fn as_ref(&self) -> &N { + &self.0 + } +} + /// Non blocking reader /// /// Typically you'd need that to check for output of a process without blocking your thread. @@ -200,36 +291,39 @@ impl NBReader { /// /// ``` /// # use std::io::Cursor; - /// use rexpect::reader::{NBReader, ReadUntil, Regex}; + /// use rexpect::reader::{NBReader, Regex, EOF, NBytes, Regx, Str}; /// // 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 (first_line, _) = e.read_until(&ReadUntil::String('\n'.to_string())).unwrap(); + /// let first_line = e.read_until(&Str("\n")).unwrap(); /// assert_eq!("Hello, miss!", &first_line); /// - /// let (_, two_bytes) = e.read_until(&ReadUntil::NBytes(2)).unwrap(); + /// let two_bytes = e.read_until(&NBytes(2)).unwrap(); /// assert_eq!("Wh", &two_bytes); /// /// let re = Regex::new(r"'[a-z]+'").unwrap(); // will find 'miss' - /// let (before, reg_match) = e.read_until(&ReadUntil::Regex(re)).unwrap(); + /// let (before, reg_match) = e.read_until(&Regx(re)).unwrap(); /// assert_eq!("at do you mean: ", &before); /// assert_eq!("'miss'", ®_match); /// - /// let (_, until_end) = e.read_until(&ReadUntil::EOF).unwrap(); + /// let until_end = e.read_until(&EOF).unwrap(); /// assert_eq!("?", &until_end); /// ``` /// - pub fn read_until(&mut self, needle: &ReadUntil) -> Result<(String, String)> { + pub fn read_until(&mut self, needle: &N) -> Result + where + N: Needle + fmt::Display + ?Sized + { let start = time::Instant::now(); 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(); - return Ok((first, second)); + if let Some(m) = needle.find(&self.buffer, self.eof) { + self.buffer.drain(..m.begin); + self.buffer.drain(..m.end - m.begin); + return Ok(m.interest); } // reached end of stream and didn't match -> error @@ -277,11 +371,9 @@ mod tests { fn test_expect_melon() { let f = io::Cursor::new("a melon\r\n"); let mut r = NBReader::new(f, None); - assert_eq!(("a melon".to_string(), "\r\n".to_string()), - r.read_until(&ReadUntil::String("\r\n".to_string())) - .expect("cannot read line")); + assert_eq!("a melon".to_owned(), r.read_until(&Str("\r\n")).expect("cannot read line")); // check for EOF - match r.read_until(&ReadUntil::NBytes(10)) { + match r.read_until(&NBytes(10)) { Ok(_) => assert!(false), Err(Error(ErrorKind::EOF(_, _, _), _)) => {} Err(Error(_, _)) => assert!(false), @@ -293,8 +385,7 @@ mod tests { let f = io::Cursor::new("2014-03-15"); let mut r = NBReader::new(f, None); let re = Regex::new(r"^\d{4}-\d{2}-\d{2}$").unwrap(); - r.read_until(&ReadUntil::Regex(re)) - .expect("regex doesn't match"); + r.read_until(&Regx(re)).expect("regex doesn't match"); } #[test] @@ -303,36 +394,32 @@ mod tests { let mut r = NBReader::new(f, None); let re = Regex::new(r"-\d{2}-").unwrap(); assert_eq!(("2014".to_string(), "-03-".to_string()), - r.read_until(&ReadUntil::Regex(re)) - .expect("regex doesn't match")); + r.read_until(&Regx(re)).expect("regex doesn't match")); } #[test] fn test_nbytes() { let f = io::Cursor::new("abcdef"); let mut r = NBReader::new(f, None); - assert_eq!(("".to_string(), "ab".to_string()), - r.read_until(&ReadUntil::NBytes(2)).expect("2 bytes")); - assert_eq!(("".to_string(), "cde".to_string()), - r.read_until(&ReadUntil::NBytes(3)).expect("3 bytes")); - assert_eq!(("".to_string(), "f".to_string()), - r.read_until(&ReadUntil::NBytes(4)).expect("4 bytes")); + assert_eq!("ab".to_string(), r.read_until(&NBytes(2)).expect("2 bytes")); + assert_eq!("cde".to_string(), r.read_until(&NBytes(3)).expect("3 bytes")); + assert_eq!("f".to_string(), r.read_until(&NBytes(4)).expect("4 bytes")); } #[test] fn test_eof() { let f = io::Cursor::new("lorem ipsum dolor sit amet"); let mut r = NBReader::new(f, None); - 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")); + r.read_until(&NBytes(2)).expect("2 bytes"); + assert_eq!("rem ipsum dolor sit amet".to_string(), + r.read_until(&EOF).expect("reading until EOF")); } #[test] fn test_try_read() { let f = io::Cursor::new("lorem"); let mut r = NBReader::new(f, None); - r.read_until(&ReadUntil::NBytes(4)).expect("4 bytes"); + r.read_until(&NBytes(4)).expect("4 bytes"); assert_eq!(Some('m'), r.try_read()); assert_eq!(None, r.try_read()); assert_eq!(None, r.try_read()); diff --git a/src/session.rs b/src/session.rs index e01b3a3c..53db4635 100644 --- a/src/session.rs +++ b/src/session.rs @@ -1,8 +1,7 @@ //! Main module of rexpect: start new process and interact with it use crate::process::PtyProcess; -use crate::reader::{NBReader, Regex}; -pub use crate::reader::ReadUntil; +use crate::reader::{NBReader, Regex, EOF, Needle, Str, Regx}; use std::fs::File; use std::io::LineWriter; use std::process::Command; @@ -89,7 +88,7 @@ impl PtySession { } // wrapper around reader::read_until to give more context for errors - fn exp(&mut self, needle: &ReadUntil) -> Result<(String, String)> { + pub fn exp(&mut self, needle: &N) -> Result { match self.reader.read_until(needle) { Ok(s) => Ok(s), Err(Error(ErrorKind::EOF(expected, got, _), _)) => { @@ -107,8 +106,8 @@ impl PtySession { /// Read one line (blocking!) and return line without the newline /// (waits until \n is in the output fetches the line and removes \r at the end if present) pub fn read_line(&mut self) -> Result { - match self.exp(&ReadUntil::String('\n'.to_string())) { - Ok((mut line, _)) => { + match self.exp(&Str("\n")) { + Ok(mut line) => { if line.ends_with('\r') { line.pop().expect("this never happens"); } @@ -127,7 +126,7 @@ impl PtySession { /// Wait until we see EOF (i.e. child process has terminated) /// Return all the yet unread output pub fn exp_eof(&mut self) -> Result { - self.exp(&ReadUntil::EOF).and_then(|(_, s)| Ok(s)) + self.exp(&EOF).and_then(|s| Ok(s)) } /// Wait until provided regex is seen on stdout of child process. @@ -138,7 +137,7 @@ impl PtySession { /// Note that `exp_regex("^foo")` matches the start of the yet consumed output. /// For matching the start of the line use `exp_regex("\nfoo")` pub fn exp_regex(&mut self, regex: &str) -> Result<(String, String)> { - let res = self.exp(&ReadUntil::Regex(Regex::new(regex).chain_err(|| "invalid regex")?)) + let res = self.exp(&Regx(Regex::new(regex).chain_err(|| "invalid regex")?)) .and_then(|s| Ok(s)); res } @@ -146,41 +145,13 @@ impl PtySession { /// Wait until provided string is seen on stdout of child process. /// Return the yet unread output (without the matched string) pub fn exp_string(&mut self, needle: &str) -> Result { - self.exp(&ReadUntil::String(needle.to_string())) - .and_then(|(s, _)| Ok(s)) + self.exp(&Str(needle)) } /// Wait until provided char is seen on stdout of child process. /// Return the yet unread output (without the matched char) pub fn exp_char(&mut self, needle: char) -> Result { - self.exp(&ReadUntil::String(needle.to_string())) - .and_then(|(s, _)| Ok(s)) - } - - /// Wait until any of the provided needles is found. - /// - /// Return a tuple with: - /// 1. the yet unread string, without the matching needle (empty in case of EOF and NBytes) - /// 2. the matched string - /// - /// # Example: - /// - /// ``` - /// use rexpect::{spawn, ReadUntil}; - /// # use rexpect::errors::*; - /// - /// # fn main() { - /// # || -> Result<()> { - /// let mut s = spawn("cat", Some(1000))?; - /// s.send_line("hello, polly!")?; - /// s.exp_any(vec![ReadUntil::String("hello".into()), - /// ReadUntil::EOF])?; - /// # Ok(()) - /// # }().expect("test failed"); - /// # } - /// ``` - pub fn exp_any(&mut self, needles: Vec) -> Result<(String, String)> { - self.exp(&ReadUntil::Any(needles)) + self.exp(&Str(needle.to_string())) } } @@ -417,6 +388,7 @@ pub fn spawn_python(timeout: Option) -> Result { #[cfg(test)] mod tests { use super::*; + use super::super::reader::{NBytes, Until, OrInterest}; #[test] fn test_read_line() { @@ -484,8 +456,31 @@ mod tests { || -> Result<()> { let mut p = spawn("cat", Some(1000)).expect("cannot run cat"); p.send_line("Hi")?; - match p.exp_any(vec![ReadUntil::NBytes(3), ReadUntil::String("Hi".to_string())]) { - Ok(s) => assert_eq!(("".to_string(), "Hi\r".to_string()), s), + + let until = Until(NBytes(3)).or(Str("Hi")); + + match p.exp(until.as_ref()) { + Ok(OrInterest::Lhs(s)) => assert_eq!("Hi\r".to_string(), s), + Ok(OrInterest::Rhs(_)) => assert!(false), + Err(e) => assert!(false, format!("got error: {}", e)), + } + Ok(()) + }() + .unwrap_or_else(|e| panic!("test_expect_any failed: {}", e)); + } + + #[test] + fn test_expect_any_huge() { + || -> Result<()> { + let mut p = spawn("cat", Some(1000)).expect("cannot run cat"); + p.send_line("Hello World")?; + + let until = Until(Str("Hi")).or(Str("World")).or(NBytes(3)); + + match p.exp(until.as_ref()) { + Ok(OrInterest::Lhs(OrInterest::Lhs(_))) => assert!(false), + Ok(OrInterest::Lhs(OrInterest::Rhs(s))) => assert_eq!("Hello ".to_string(), s), + Ok(OrInterest::Rhs(_)) => assert!(false), Err(e) => assert!(false, format!("got error: {}", e)), } Ok(()) From d8b0f20e500a1984a25cbdf80789960251400fcc Mon Sep 17 00:00:00 2001 From: Maxim Zhiburt Date: Mon, 8 Jun 2020 00:36:34 +0300 Subject: [PATCH 4/4] resolve merge conficts Signed-off-by: Maxim Zhiburt --- src/session.rs | 27 +++++---------------------- 1 file changed, 5 insertions(+), 22 deletions(-) diff --git a/src/session.rs b/src/session.rs index 0b87e1e0..25674d5b 100644 --- a/src/session.rs +++ b/src/session.rs @@ -73,13 +73,7 @@ impl StreamSession { // wrapper around reader::read_until to give more context for errors pub fn exp(&mut self, needle: &N) -> Result { - match self.reader.read_until(needle) { - Ok(s) => Ok(s), - Err(Error(ErrorKind::EOF(expected, got, _), _)) => { - Err(ErrorKind::EOF(expected, got, self.process.status()).into()) - } - Err(e) => Err(e), - } + self.reader.read_until(needle) } /// Make sure all bytes written via `send()` are sent to the process @@ -107,11 +101,6 @@ impl StreamSession { self.reader.try_read() } - // wrapper around reader::read_until to give more context for errors - fn exp(&mut self, needle: &ReadUntil) -> Result<(String, String)> { - self.reader.read_until(needle) - } - /// Wait until we see EOF (i.e. child process has terminated) /// Return all the yet unread output pub fn exp_eof(&mut self) -> Result { @@ -238,18 +227,11 @@ pub fn spawn(program: &str, timeout_ms: Option) -> Result { /// See `spawn` pub fn spawn_command(command: Command, timeout_ms: Option) -> Result { let commandname = format!("{:?}", &command); - let mut process = PtyProcess::new(command).chain_err(|| "couldn't start process")?; + let mut process = PtyProcess::new(command) + .chain_err(|| "couldn't start process")?; process.set_kill_timeout(timeout_ms); - let f = process.get_file_handle(); - let writer = LineWriter::new(f.try_clone().chain_err(|| "couldn't open write stream")?); - let reader = NBReader::new(f, timeout_ms); - Ok(PtySession { - process: process, - writer: writer, - reader: reader, - commandname: commandname, - }) + PtySession::new(process, timeout_ms, commandname) } /// A repl session: e.g. bash or the python shell: @@ -545,6 +527,7 @@ mod tests { Ok(()) }() .unwrap_or_else(|e| panic!("test_expect_any failed: {}", e)); + } #[test] fn test_expect_empty_command_error() {