Merge pull request #69 from takluyver/asyncio

asyncio integration for expect

6 files changed, 283 insertions, 80 deletions

diff --git a/.travis.yml b/.travis.yml
index 4a68168..3a2f331 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,7 +1,6 @@
 language: python
 
 python:
-  - 2.6
   - 2.7
   - 3.3
   - 3.4
diff --git a/doc/history.rst b/doc/history.rst
index a09f124..ec1c9c3 100644
--- a/doc/history.rst
+++ b/doc/history.rst
@@ -4,6 +4,15 @@ History
 Releases
 --------
 
+Version 4.0
+```````````
+
+* Integration with :mod:`asyncio`: passing ``async=True`` to :meth:`~.expect`,
+  :meth:`~.expect_exact` or :meth:`~.expect_list` will make them return a
+  coroutine. You can get the result using ``yield from``, or wrap it in an
+  :class:`asyncio.Task`. This allows the event loop to do other things while
+  waiting for output that matches a pattern.
+
 Version 3.4
 ```````````
 * Fixed regression when executing pexpect with some prior releases of
diff --git a/pexpect/__init__.py b/pexpect/__init__.py
index 06dbac4..f070067 100644
--- a/pexpect/__init__.py
+++ b/pexpect/__init__.py
@@ -88,6 +88,8 @@ except ImportError:  # pragma: no cover
 A critical module was not found. Probably this operating system does not
 support it. Pexpect is intended for UNIX-like operating systems.''')
 
+from .expect import Expecter
+
 __version__ = '3.3'
 __revision__ = ''
 __all__ = ['ExceptionPexpect', 'EOF', 'TIMEOUT', 'spawn', 'spawnu', 'run', 'runu',
@@ -113,7 +115,8 @@ class ExceptionPexpect(Exception):
         is not included. '''
 
         tblist = traceback.extract_tb(sys.exc_info()[2])
-        tblist = [item for item in tblist if 'pexpect/__init__' not in item[0]]
+        tblist = [item for item in tblist if ('pexpect/__init__' not in item[0])
+                                           and ('pexpect/expect' not in item[0])]
         tblist = traceback.format_list(tblist)
         return ''.join(tblist)
 
@@ -1382,7 +1385,7 @@ class spawn(object):
                 self._pattern_type_err(p)
         return compiled_pattern_list
 
-    def expect(self, pattern, timeout=-1, searchwindowsize=-1):
+    def expect(self, pattern, timeout=-1, searchwindowsize=-1, async=False):
 
         '''This seeks through the stream until a pattern is matched. The
         pattern is overloaded and may take several types. The pattern can be a
@@ -1457,14 +1460,25 @@ class spawn(object):
                 print p.before
 
         If you are trying to optimize for speed then see expect_list().
+        
+        On Python 3.4, or Python 3.3 with asyncio installed, passing
+        ``async=True``  will make this return an :mod:`asyncio` coroutine,
+        which you can yield from to get the same result that this method would
+        normally give directly. So, inside a coroutine, you can replace this code::
+        
+            index = p.expect(patterns)
+        
+        With this non-blocking form::
+        
+            index = yield from p.expect(patterns, async=True)
         '''
 
         compiled_pattern_list = self.compile_pattern_list(pattern)
         return self.expect_list(compiled_pattern_list,
-                timeout, searchwindowsize)
-
-    def expect_list(self, pattern_list, timeout=-1, searchwindowsize=-1):
+                timeout, searchwindowsize, async)
 
+    def expect_list(self, pattern_list, timeout=-1, searchwindowsize=-1,
+                    async=False):
         '''This takes a list of compiled regular expressions and returns the
         index into the pattern_list that matched the child output. The list may
         also contain EOF or TIMEOUT(which are not compiled regular
@@ -1473,12 +1487,23 @@ class spawn(object):
         may help if you are trying to optimize for speed, otherwise just use
         the expect() method.  This is called by expect(). If timeout==-1 then
         the self.timeout value is used. If searchwindowsize==-1 then the
-        self.searchwindowsize value is used. '''
+        self.searchwindowsize value is used.
+
+        Like :meth:`expect`, passing ``async=True`` will make this return an
+        asyncio coroutine.
+        '''
+        if timeout == -1:
+            timeout = self.timeout
 
-        return self.expect_loop(searcher_re(pattern_list),
-                timeout, searchwindowsize)
+        exp = Expecter(self, searcher_re(pattern_list), searchwindowsize)
+        if async:
+            from .async import expect_async
+            return expect_async(exp, timeout)
+        else:
+            return exp.expect_loop(timeout)
 
-    def expect_exact(self, pattern_list, timeout=-1, searchwindowsize=-1):
+    def expect_exact(self, pattern_list, timeout=-1, searchwindowsize=-1,
+                     async=False):
 
         '''This is similar to expect(), but uses plain string matching instead
         of compiled regular expressions in 'pattern_list'. The 'pattern_list'
@@ -1490,7 +1515,13 @@ class spawn(object):
         search to just the end of the input buffer.
 
         This method is also useful when you don't want to have to worry about
-        escaping regular expression characters that you want to match.'''
+        escaping regular expression characters that you want to match.
+        
+        Like :meth:`expect`, passing ``async=True`` will make this return an
+        asyncio coroutine.
+        '''
+        if timeout == -1:
+            timeout = self.timeout
 
         if (isinstance(pattern_list, self.allowed_string_types) or
                 pattern_list in (TIMEOUT, EOF)):
@@ -1508,83 +1539,23 @@ class spawn(object):
         except TypeError:
             self._pattern_type_err(pattern_list)
         pattern_list = [prepare_pattern(p) for p in pattern_list]
-        return self.expect_loop(searcher_string(pattern_list),
-                timeout, searchwindowsize)
 
-    def expect_loop(self, searcher, timeout=-1, searchwindowsize=-1):
+        exp = Expecter(self, searcher_string(pattern_list), searchwindowsize)
+        if async:
+            from .async import expect_async
+            return expect_async(exp, timeout)
+        else:
+            return exp.expect_loop(timeout)
 
+    def expect_loop(self, searcher, timeout=-1, searchwindowsize=-1):
         '''This is the common loop used inside expect. The 'searcher' should be
         an instance of searcher_re or searcher_string, which describes how and
         what to search for in the input.
 
         See expect() for other arguments, return value and exceptions. '''
 
-        self.searcher = searcher
-
-        if timeout == -1:
-            timeout = self.timeout
-        if timeout is not None:
-            end_time = time.time() + timeout
-        if searchwindowsize == -1:
-            searchwindowsize = self.searchwindowsize
-
-        try:
-            incoming = self.buffer
-            freshlen = len(incoming)
-            while True:
-                # Keep reading until exception or return.
-                index = searcher.search(incoming, freshlen, searchwindowsize)
-                if index >= 0:
-                    self.buffer = incoming[searcher.end:]
-                    self.before = incoming[: searcher.start]
-                    self.after = incoming[searcher.start: searcher.end]
-                    self.match = searcher.match
-                    self.match_index = index
-                    return self.match_index
-                # No match at this point
-                if (timeout is not None) and (timeout < 0):
-                    raise TIMEOUT('Timeout exceeded in expect_any().')
-                # Still have time left, so read more data
-                c = self.read_nonblocking(self.maxread, timeout)
-                freshlen = len(c)
-                time.sleep(0.0001)
-                incoming = incoming + c
-                if timeout is not None:
-                    timeout = end_time - time.time()
-        except EOF:
-            err = sys.exc_info()[1]
-            self.buffer = self.string_type()
-            self.before = incoming
-            self.after = EOF
-            index = searcher.eof_index
-            if index >= 0:
-                self.match = EOF
-                self.match_index = index
-                return self.match_index
-            else:
-                self.match = None
-                self.match_index = None
-                raise EOF(str(err) + '\n' + str(self))
-        except TIMEOUT:
-            err = sys.exc_info()[1]
-            self.buffer = incoming
-            self.before = incoming
-            self.after = TIMEOUT
-            index = searcher.timeout_index
-            if index >= 0:
-                self.match = TIMEOUT
-                self.match_index = index
-                return self.match_index
-            else:
-                self.match = None
-                self.match_index = None
-                raise TIMEOUT(str(err) + '\n' + str(self))
-        except:
-            self.before = incoming
-            self.after = None
-            self.match = None
-            self.match_index = None
-            raise
+        exp = Expecter(self, searcher, searchwindowsize)
+        return exp.expect_loop(timeout)
 
     def getwinsize(self):
 
diff --git a/pexpect/async.py b/pexpect/async.py
new file mode 100644
index 0000000..8ec9c3c
--- /dev/null
+++ b/pexpect/async.py
@@ -0,0 +1,68 @@
+import asyncio
+import errno
+
+from pexpect import EOF
+
+@asyncio.coroutine
+def expect_async(expecter, timeout=None):
+    # First process data that was previously read - if it maches, we don't need
+    # async stuff.    
+    idx = expecter.new_data(expecter.spawn.buffer)
+    expecter.spawn.buffer = expecter.spawn.string_type()
+    if idx:
+        return idx
+
+    transport, pw = yield from asyncio.get_event_loop()\
+        .connect_read_pipe(lambda: PatternWaiter(expecter), expecter.spawn)
+
+    try:
+        return (yield from asyncio.wait_for(pw.fut, timeout))
+    except asyncio.TimeoutError as e:
+        transport.pause_reading()
+        return expecter.timeout(e)
+
+class PatternWaiter(asyncio.Protocol):
+    def __init__(self, expecter):
+        self.expecter = expecter
+        self.fut = asyncio.Future()
+    
+    def found(self, result):
+        if not self.fut.done():
+            self.fut.set_result(result)
+    
+    def error(self, exc):
+        if not self.fut.done():
+            self.fut.set_exception(exc)
+    
+    def data_received(self, data):
+        spawn = self.expecter.spawn
+        s = spawn._coerce_read_string(data)
+        spawn._log(s, 'read')
+
+        if self.fut.done():
+            spawn.buffer += data
+            return
+
+        try:
+            index = self.expecter.new_data(data)
+            if index is not None:
+                # Found a match
+                self.found(index)
+        except Exception as e:
+            self.expecter.errored()
+            self.error(e)
+    
+    def eof_received(self):
+        try:
+            index = self.expecter.eof()
+        except EOF as e:
+            self.error(e)
+        else:
+            self.found(index)
+    
+    def connection_lost(self, exc):
+        if isinstance(exc, OSError) and exc.errno == errno.EIO:
+            # We may get here without eof_received being called, e.g on Linux
+            self.eof_received()
+        elif exc is not None:
+            self.error(exc)
\ No newline at end of file
diff --git a/pexpect/expect.py b/pexpect/expect.py
new file mode 100644
index 0000000..b8da406
--- /dev/null
+++ b/pexpect/expect.py
@@ -0,0 +1,105 @@
+import time
+
+class Expecter(object):
+    def __init__(self, spawn, searcher, searchwindowsize=-1):
+        self.spawn = spawn
+        self.searcher = searcher
+        if searchwindowsize == -1:
+            searchwindowsize = spawn.searchwindowsize
+        self.searchwindowsize = searchwindowsize
+    
+    def new_data(self, data):
+        spawn = self.spawn
+        searcher = self.searcher
+
+        incoming = spawn.buffer + data
+        freshlen = len(data)
+        index = searcher.search(incoming, freshlen, self.searchwindowsize)
+        if index >= 0:
+            spawn.buffer = incoming[searcher.end:]
+            spawn.before = incoming[: searcher.start]
+            spawn.after = incoming[searcher.start: searcher.end]
+            spawn.match = searcher.match
+            spawn.match_index = index
+            # Found a match
+            return index
+    
+        spawn.buffer = incoming
+    
+    def eof(self, err=None):
+        spawn = self.spawn
+        from . import EOF
+
+        spawn.before = spawn.buffer
+        spawn.buffer = spawn.string_type()
+        spawn.after = EOF
+        index = self.searcher.eof_index
+        if index >= 0:
+            spawn.match = EOF
+            spawn.match_index = index
+            return index
+        else:
+            spawn.match = None
+            spawn.match_index = None
+            msg = str(spawn)
+            if err is not None:
+                msg = str(err) + '\n' + msg
+            raise EOF(msg)
+    
+    def timeout(self, err=None):
+        spawn = self.spawn
+        from . import TIMEOUT
+
+        spawn.before = spawn.buffer
+        spawn.after = TIMEOUT
+        index = self.searcher.timeout_index
+        if index >= 0:
+            spawn.match = TIMEOUT
+            spawn.match_index = index
+            return index
+        else:
+            spawn.match = None
+            spawn.match_index = None
+            msg = str(spawn)
+            if err is not None:
+                msg = str(err) + '\n' + msg
+            raise TIMEOUT(msg)
+
+    def errored(self):
+        spawn = self.spawn
+        spawn.before = spawn.buffer
+        spawn.after = None
+        spawn.match = None
+        spawn.match_index = None
+    
+    def expect_loop(self, timeout=-1):
+        """Blocking expect"""
+        spawn = self.spawn
+        from . import EOF, TIMEOUT
+
+        if timeout is not None:
+            end_time = time.time() + timeout
+
+        try:
+            incoming = spawn.buffer
+            spawn.buffer = spawn.string_type()  # Treat buffer as new data
+            while True:
+                idx = self.new_data(incoming)
+                # Keep reading until exception or return.
+                if idx is not None:
+                    return idx
+                # No match at this point
+                if (timeout is not None) and (timeout < 0):
+                    return self.timeout()
+                # Still have time left, so read more data
+                incoming = spawn.read_nonblocking(spawn.maxread, timeout)
+                time.sleep(0.0001)
+                if timeout is not None:
+                    timeout = end_time - time.time()
+        except EOF as e:
+            return self.eof(e)
+        except TIMEOUT as e:
+            return self.timeout(e)
+        except:
+            self.errored()
+            raise
\ No newline at end of file
diff --git a/tests/test_async.py b/tests/test_async.py
new file mode 100644
index 0000000..ce75572
--- /dev/null
+++ b/tests/test_async.py
@@ -0,0 +1,51 @@
+try:
+    import asyncio
+except ImportError:
+    asyncio = None
+
+import sys
+import unittest
+
+import pexpect
+from .PexpectTestCase import PexpectTestCase
+
+def run(coro):
+    return asyncio.get_event_loop().run_until_complete(coro)
+
+@unittest.skipIf(asyncio is None, "Requires asyncio")
+class AsyncTests(PexpectTestCase):
+    def test_simple_expect(self):
+        p = pexpect.spawn('cat')
+        p.sendline('Hello asyncio')
+        coro = p.expect(['Hello', pexpect.EOF] , async=True)
+        assert run(coro) == 0
+        print('Done')
+
+    def test_timeout(self):
+        p = pexpect.spawn('cat')
+        coro = p.expect('foo', timeout=1, async=True)
+        with self.assertRaises(pexpect.TIMEOUT):
+            run(coro)
+        
+        p = pexpect.spawn('cat')
+        coro = p.expect(['foo', pexpect.TIMEOUT], timeout=1, async=True)
+        assert run(coro) == 1
+
+    def test_eof(self):
+        p = pexpect.spawn('cat')
+        p.sendline('Hi')
+        coro = p.expect(pexpect.EOF, async=True)
+        p.sendeof()
+        assert run(coro) == 0
+
+        p = pexpect.spawn('cat')
+        p.sendeof()
+        coro = p.expect('Blah', async=True)
+        with self.assertRaises(pexpect.EOF):
+            run(coro)
+    
+    def test_expect_exact(self):
+        p = pexpect.spawn('%s list100.py' % sys.executable)
+        assert run(p.expect_exact(b'5', async=True)) == 0
+        assert run(p.expect_exact(['wpeok', b'11'], async=True)) == 1
+        assert run(p.expect_exact([b'foo', pexpect.EOF], async=True)) == 1