make match_any also return index of matched needle

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 description = "Interact with unix processes/bash the same way as pexpect or Don libes expect does"
 name        = "rexpect"
-version     = "0.3.0"
+version     = "0.4.0"
 authors     = ["Philipp Keller <[email protected]>"]
 edition     = "2018"
 repository  = "https://github.com/philippkeller/rexpect"
@@ -20,4 +20,4 @@ error-chain = "0.12"
 tempfile = "3"
 
 [badges]
-travis-ci = { repository = "philippkeller/rexpect" }
+travis-ci = { repository = "philippkeller/rexpect" }

src/process.rs

@@ -71,7 +71,6 @@ use nix::pty::ptsname_r;
 /// based on https://blog.tarq.io/ptsname-on-osx-with-rust/
 fn ptsname_r(fd: &PtyMaster) -> nix::Result<String> {
     use std::ffi::CStr;
-    use std::os::unix::io::AsRawFd;
     use nix::libc::{ioctl, TIOCPTYGNAME};
 
     // the buffer size on OSX is 128, defined by sys/ttycom.h

src/reader.rs

@@ -93,6 +93,15 @@ pub fn find(needle: &ReadUntil, buffer: &str, eof: bool) -> Option<(usize, usize
     }
 }
 
+pub fn find_any(needle: &Vec<ReadUntil>, buffer: &str, eof: bool) -> Option<(usize, usize, usize)> {
+    for (pos, read_until) in needle.iter().enumerate() {
+        if let Some((before, after)) = find(&read_until, buffer, eof) {
+            return Some((before, after, pos));
+        }
+    }
+    None
+}
+
 /// Non blocking reader
 ///
 /// Typically you'd need that to check for output of a process without blocking your thread.
@@ -222,14 +231,55 @@ impl NBReader {
     /// ```
     ///
     pub fn read_until(&mut self, needle: &ReadUntil) -> Result<(String, String)> {
+        match self.read_until_tuple_pos(needle) {
+            Ok(tuple_pos) => {
+                let first = self.buffer.drain(..tuple_pos.0).collect();
+                let second = self.buffer.drain(..tuple_pos.1 - tuple_pos.0).collect();
+                Ok((first, second))
+            },
+            Err(e) => Err(e),
+        }
+    }
+
+    /// Read until any needle is found (blocking!) and return tuple with:
+    /// 1. yet unread string until and without needle
+    /// 2. matched needle
+    /// 3. index of which needle is found
+    /// 
+    /// this is the same function as read_until with the difference that the index
+    /// is also returned
+    pub fn read_until_any(&mut self, needle: Vec<ReadUntil>) -> Result<(String, String, usize)> {
+        match self.read_until_tuple_pos(&ReadUntil::Any(needle)) {
+            Ok(tuple_pos) => {
+                let first = self.buffer.drain(..tuple_pos.0).collect();
+                let second = self.buffer.drain(..tuple_pos.1 - tuple_pos.0).collect();
+                Ok((first, second, tuple_pos.2))
+            },
+            Err(e) => Err(e),
+        }
+    }
+    /// Read until needle is found and return
+    /// 1. yet unread string until and without needle
+    /// 2. matched needle
+    /// 3. index of which needle is found. This is only making sense for ReadUntil::Any
+    ///    for the other types this is redundant but I didn't find any more elegant way
+    ///    to solve this.
+    fn read_until_tuple_pos(&mut self, needle: &ReadUntil) -> Result<(usize, usize, usize)> {
         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));
+            match needle {
+                ReadUntil::Any(n) => {
+                    if let Some(tuple_pos) = find_any(n, &self.buffer, self.eof) {
+                        return Ok(tuple_pos)
+                    }
+                }
+                needle => {
+                    if let Some((before, after)) = find(needle, &self.buffer, self.eof) {
+                        return Ok((before, after, 0))
+                    }
+                }
             }
 
             // reached end of stream and didn't match -> error

src/session.rs

@@ -99,6 +99,17 @@ impl PtySession {
         }
     }
 
+    // wrapper around reader::read_until to give more context for errors
+    fn exp_any_wrapper(&mut self, needle: Vec<ReadUntil>) -> Result<(String, String, usize)> {
+        match self.reader.read_until_any(needle) {
+            Ok(s) => Ok(s),
+            Err(Error(ErrorKind::EOF(expected, got, _), _)) => {
+                Err(ErrorKind::EOF(expected, got, self.process.status()).into())
+            }
+            Err(e) => Err(e),
+        }
+    }
+
     /// Make sure all bytes written via `send()` are sent to the process
     pub fn flush(&mut self) -> Result<()> {
         self.writer.flush().chain_err(|| "could not flush")
@@ -179,8 +190,8 @@ impl PtySession {
     ///     # }().expect("test failed");
     /// # }
     /// ```
-    pub fn exp_any(&mut self, needles: Vec<ReadUntil>) -> Result<(String, String)> {
-        self.exp(&ReadUntil::Any(needles))
+    pub fn exp_any(&mut self, needles: Vec<ReadUntil>) -> Result<(String, String, usize)> {
+        self.exp_any_wrapper(needles)
     }
 }
 
@@ -483,8 +494,8 @@ 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),
+            match p.exp_any(vec![ReadUntil::String("Bye".to_string()), ReadUntil::String("Hi".to_string())]) {
+                Ok(s) => assert_eq!(("".to_string(), "Hi".to_string(), 1), s),
                 Err(e) => assert!(false, format!("got error: {}", e)),
             }
             Ok(())