Allow for configurable select.poll() usage

- This allows systems that can have > 1024 fds to work
- Should be marginally faster
- Added tests to use select.poll()
- Incremented version - Happy to change this, guessed what it should be

9 files changed, 117 insertions, 23 deletions

diff --git a/.gitignore b/.gitignore
index 914c903..22cd478 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,3 +8,4 @@ MANIFEST
 .coverage*
 htmlcov
 *.egg-info/
+.cache/
diff --git a/pexpect/__init__.py b/pexpect/__init__.py
index 458d3ad..83462a4 100644
--- a/pexpect/__init__.py
+++ b/pexpect/__init__.py
@@ -75,7 +75,7 @@ if sys.platform != 'win32':
     from .pty_spawn import spawn, spawnu
     from .run import run, runu
 
-__version__ = '4.4.0'
+__version__ = '4.5.0'
 __revision__ = ''
 __all__ = ['ExceptionPexpect', 'EOF', 'TIMEOUT', 'spawn', 'spawnu', 'run', 'runu',
            'which', 'split_command_line', '__version__', '__revision__']
diff --git a/pexpect/fdpexpect.py b/pexpect/fdpexpect.py
index cd60804..cddd50e 100644
--- a/pexpect/fdpexpect.py
+++ b/pexpect/fdpexpect.py
@@ -23,7 +23,7 @@ PEXPECT LICENSE
 
 from .spawnbase import SpawnBase
 from .exceptions import ExceptionPexpect, TIMEOUT
-from .utils import select_ignore_interrupts
+from .utils import select_ignore_interrupts, poll_ignore_interrupts
 import os
 
 __all__ = ['fdspawn']
@@ -34,7 +34,7 @@ class fdspawn(SpawnBase):
     for patterns, or to control a modem or serial device. '''
 
     def __init__ (self, fd, args=None, timeout=30, maxread=2000, searchwindowsize=None,
-                  logfile=None, encoding=None, codec_errors='strict'):
+                  logfile=None, encoding=None, codec_errors='strict', use_poll=False):
         '''This takes a file descriptor (an int) or an object that support the
         fileno() method (returning an int). All Python file-like objects
         support fileno(). '''
@@ -58,6 +58,7 @@ class fdspawn(SpawnBase):
         self.own_fd = False
         self.closed = False
         self.name = '<file descriptor %d>' % fd
+        self.use_poll = use_poll
 
     def close (self):
         """Close the file descriptor.
@@ -88,7 +89,7 @@ class fdspawn(SpawnBase):
     def terminate (self, force=False):  # pragma: no cover
         '''Deprecated and invalid. Just raises an exception.'''
         raise ExceptionPexpect('This method is not valid for file descriptors.')
-    
+
     # These four methods are left around for backwards compatibility, but not
     # documented as part of fdpexpect. You're encouraged to use os.write
     # directly.
@@ -96,19 +97,19 @@ class fdspawn(SpawnBase):
         "Write to fd, return number of bytes written"
         s = self._coerce_send_string(s)
         self._log(s, 'send')
-        
+
         b = self._encoder.encode(s, final=False)
         return os.write(self.child_fd, b)
-    
+
     def sendline(self, s):
         "Write to fd with trailing newline, return number of bytes written"
         s = self._coerce_send_string(s)
         return self.send(s + self.linesep)
-    
+
     def write(self, s):
         "Write to fd, return None"
         self.send(s)
-    
+
     def writelines(self, sequence):
         "Call self.write() for each item in sequence"
         for s in sequence:
@@ -136,7 +137,12 @@ class fdspawn(SpawnBase):
             rlist = [self.child_fd]
             wlist = []
             xlist = []
-            rlist, wlist, xlist = select_ignore_interrupts(rlist, wlist, xlist, timeout)
+            if self.use_poll:
+                rlist = poll_ignore_interrupts(rlist, timeout)
+            else:
+                rlist, wlist, xlist = select_ignore_interrupts(
+                    rlist, wlist, xlist, timeout
+                )
             if self.child_fd not in rlist:
                 raise TIMEOUT('Timeout exceeded.')
         return super(fdspawn, self).read_nonblocking(size)
diff --git a/pexpect/pty_spawn.py b/pexpect/pty_spawn.py
index c1a835b..403689b 100644
--- a/pexpect/pty_spawn.py
+++ b/pexpect/pty_spawn.py
@@ -12,7 +12,9 @@ from ptyprocess.ptyprocess import use_native_pty_fork
 
 from .exceptions import ExceptionPexpect, EOF, TIMEOUT
 from .spawnbase import SpawnBase
-from .utils import which, split_command_line, select_ignore_interrupts
+from .utils import (
+    which, split_command_line, select_ignore_interrupts, poll_ignore_interrupts
+)
 
 @contextmanager
 def _wrap_ptyprocess_err():
@@ -34,7 +36,8 @@ class spawn(SpawnBase):
     def __init__(self, command, args=[], timeout=30, maxread=2000,
                  searchwindowsize=None, logfile=None, cwd=None, env=None,
                  ignore_sighup=False, echo=True, preexec_fn=None,
-                 encoding=None, codec_errors='strict', dimensions=None):
+                 encoding=None, codec_errors='strict', dimensions=None,
+                 use_poll=False):
         '''This is the constructor. The command parameter may be a string that
         includes a command and any arguments to the command. For example::
 
@@ -171,7 +174,7 @@ class spawn(SpawnBase):
         using setecho(False) followed by waitnoecho().  However, for some
         platforms such as Solaris, this is not possible, and should be
         disabled immediately on spawn.
-        
+
         If preexec_fn is given, it will be called in the child process before
         launching the given command. This is useful to e.g. reset inherited
         signal handlers.
@@ -179,6 +182,9 @@ class spawn(SpawnBase):
         The dimensions attribute specifies the size of the pseudo-terminal as
         seen by the subprocess, and is specified as a two-entry tuple (rows,
         columns). If this is unspecified, the defaults in ptyprocess will apply.
+
+        The use_poll attribute enables using select.poll() over select.select()
+        for socket handling. This is handy if your system could have > 1024 fds
         '''
         super(spawn, self).__init__(timeout=timeout, maxread=maxread, searchwindowsize=searchwindowsize,
                                     logfile=logfile, encoding=encoding, codec_errors=codec_errors)
@@ -196,6 +202,7 @@ class spawn(SpawnBase):
             self.name = '<pexpect factory incomplete>'
         else:
             self._spawn(command, args, preexec_fn, dimensions)
+        self.use_poll = use_poll
 
     def __str__(self):
         '''This returns a human-readable string that represents the state of
@@ -439,7 +446,10 @@ class spawn(SpawnBase):
         # If isalive() is false, then I pretend that this is the same as EOF.
         if not self.isalive():
             # timeout of 0 means "poll"
-            r, w, e = select_ignore_interrupts([self.child_fd], [], [], 0)
+            if self.use_poll:
+                r = poll_ignore_interrupts([self.child_fd], timeout)
+            else:
+                r, w, e = select_ignore_interrupts([self.child_fd], [], [], 0)
             if not r:
                 self.flag_eof = True
                 raise EOF('End Of File (EOF). Braindead platform.')
@@ -447,12 +457,19 @@ class spawn(SpawnBase):
             # Irix takes a long time before it realizes a child was terminated.
             # FIXME So does this mean Irix systems are forced to always have
             # FIXME a 2 second delay when calling read_nonblocking? That sucks.
-            r, w, e = select_ignore_interrupts([self.child_fd], [], [], 2)
+            if self.use_poll:
+                r = poll_ignore_interrupts([self.child_fd], timeout)
+            else:
+                r, w, e = select_ignore_interrupts([self.child_fd], [], [], 2)
             if not r and not self.isalive():
                 self.flag_eof = True
                 raise EOF('End Of File (EOF). Slow platform.')
-
-        r, w, e = select_ignore_interrupts([self.child_fd], [], [], timeout)
+        if self.use_poll:
+            r = poll_ignore_interrupts([self.child_fd], timeout)
+        else:
+            r, w, e = select_ignore_interrupts(
+                [self.child_fd], [], [], timeout
+            )
 
         if not r:
             if not self.isalive():
@@ -771,7 +788,12 @@ class spawn(SpawnBase):
         '''
 
         while self.isalive():
-            r, w, e = select_ignore_interrupts([self.child_fd, self.STDIN_FILENO], [], [])
+            if self.use_poll:
+                r = poll_ignore_interrupts([self.child_fd], timeout)
+            else:
+                r, w, e = select_ignore_interrupts(
+                    [self.child_fd, self.STDIN_FILENO], [], []
+                )
             if self.child_fd in r:
                 try:
                     data = self.__interact_read(self.child_fd)
diff --git a/pexpect/utils.py b/pexpect/utils.py
index bafc280..117f7c9 100644
--- a/pexpect/utils.py
+++ b/pexpect/utils.py
@@ -154,3 +154,32 @@ def select_ignore_interrupts(iwtd, owtd, ewtd, timeout=None):
                 # something else caused the select.error, so
                 # this actually is an exception.
                 raise
+
+
+def poll_ignore_interrupts(fds, timeout=None):
+    """Simple wrapper around poll for a list of file descriptors."""
+
+    if timeout is not None:
+        end_time = time.time() + timeout
+
+    poller = select.poll()
+    for fd in fds:
+        poller.register(fd)
+        while True:
+            try:
+                timeout_ms = None if timeout is None else timeout * 1000
+                results = poller.poll(timeout_ms)
+                return [fd for fd, _ in results]
+            except InterruptedError:
+                err = sys.exc_info()[1]
+                if err.args[0] == errno.EINTR:
+                    # if we loop back we have to subtract the
+                    # amount of time we already waited.
+                    if timeout is not None:
+                        timeout = end_time - time.time()
+                        if timeout < 0:
+                            return []
+                else:
+                    # something else caused the select.error, so
+                    # this actually is an exception.
+                    raise
diff --git a/setup.cfg b/setup.cfg
index 87fce02..b2a82dc 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,4 +1,4 @@
-[pytest]
+[tool:pytest]
 norecursedirs = .git
 
 [bdist_wheel]
diff --git a/setup.py b/setup.py
index 2b811ee..4e61e79 100644
--- a/setup.py
+++ b/setup.py
@@ -34,7 +34,7 @@ for patterns from file descriptors or subprocesses—are also available on
 Windows.
 """
 
-setup (name='pexpect',
+setup(name='pexpect',
     version=version,
     packages=['pexpect'],
     package_data={'pexpect': ['bashrc.sh']},
diff --git a/tests/test_misc.py b/tests/test_misc.py
index fcd77d3..118de2e 100755
--- a/tests/test_misc.py
+++ b/tests/test_misc.py
@@ -52,6 +52,15 @@ class TestCaseMisc(PexpectTestCase.PexpectTestCase):
             return 'skip'
         assert child.isatty()
 
+    def test_isatty_poll(self):
+        " Test isatty() is True after spawning process on most platforms. "
+        child = pexpect.spawn('cat', use_poll=True)
+        if not child.isatty() and sys.platform.lower().startswith('sunos'):
+            if hasattr(unittest, 'SkipTest'):
+                raise unittest.SkipTest("Not supported on this platform.")
+            return 'skip'
+        assert child.isatty()
+
     def test_read(self):
         " Test spawn.read by calls of various size. "
         child = pexpect.spawn('cat')
@@ -65,6 +74,19 @@ class TestCaseMisc(PexpectTestCase.PexpectTestCase):
         remaining = child.read().replace(_CAT_EOF, b'')
         self.assertEqual(remaining, b'abc\r\n')
 
+    def test_read_poll(self):
+        " Test spawn.read by calls of various size. "
+        child = pexpect.spawn('cat', use_poll=True)
+        child.sendline("abc")
+        child.sendeof()
+        self.assertEqual(child.read(0), b'')
+        self.assertEqual(child.read(1), b'a')
+        self.assertEqual(child.read(1), b'b')
+        self.assertEqual(child.read(1), b'c')
+        self.assertEqual(child.read(2), b'\r\n')
+        remaining = child.read().replace(_CAT_EOF, b'')
+        self.assertEqual(remaining, b'abc\r\n')
+
     def test_readline_bin_echo(self):
         " Test spawn('echo'). "
         # given,
@@ -148,7 +170,7 @@ class TestCaseMisc(PexpectTestCase.PexpectTestCase):
             p.sendline(b'alpha')
             p.expect(b'<out>alpha')
             assert p.isalive()
-        
+
         assert not p.isalive()
 
     def test_terminate(self):
@@ -343,4 +365,3 @@ if __name__ == '__main__':
     unittest.main()
 
 suite = unittest.makeSuite(TestCaseMisc,'test')
-
diff --git a/tests/test_socket.py b/tests/test_socket.py
index 56d69c7..21648f4 100644
--- a/tests/test_socket.py
+++ b/tests/test_socket.py
@@ -225,7 +225,7 @@ class ExpectTestCase(PexpectTestCase.PexpectTestCase):
         session.expect(pexpect.EOF)
         self.assertEqual(session.before, b'')
 
-    def test_fd_isalive (self):
+    def test_fd_isalive(self):
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         sock.connect((self.host, self.port))
         session = fdpexpect.fdspawn(sock.fileno(), timeout=10)
@@ -233,13 +233,28 @@ class ExpectTestCase(PexpectTestCase.PexpectTestCase):
         sock.close()
         assert not session.isalive(), "Should not be alive after close()"
 
-    def test_fd_isatty (self):
+    def test_fd_isalive_poll(self):
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.connect((self.host, self.port))
+        session = fdpexpect.fdspawn(sock.fileno(), timeout=10, use_poll=True)
+        assert session.isalive()
+        sock.close()
+        assert not session.isalive(), "Should not be alive after close()"
+
+    def test_fd_isatty(self):
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         sock.connect((self.host, self.port))
         session = fdpexpect.fdspawn(sock.fileno(), timeout=10)
         assert not session.isatty()
         session.close()
 
+    def test_fd_isatty_poll(self):
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.connect((self.host, self.port))
+        session = fdpexpect.fdspawn(sock.fileno(), timeout=10, use_poll=True)
+        assert not session.isatty()
+        session.close()
+
     def test_fileobj(self):
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         sock.connect((self.host, self.port))