Merge pull request #80 from pexpect/replwrap-continuation-timeout-results

depricated: Resolve return value and timeout for multi-command and continuation prompts in replwrap

2 files changed, 40 insertions, 20 deletions

diff --git a/pexpect/replwrap.py b/pexpect/replwrap.py
index 83a09c2..6439b35 100644
--- a/pexpect/replwrap.py
+++ b/pexpect/replwrap.py
@@ -3,20 +3,16 @@
 import os.path
 import signal
 import sys
-import re
 
 import pexpect
 
 PY3 = (sys.version_info[0] >= 3)
 
 if PY3:
-    def u(s): return s
     basestring = str
-else:
-    def u(s): return s.decode('utf-8')
 
-PEXPECT_PROMPT = u('[PEXPECT_PROMPT>')
-PEXPECT_CONTINUATION_PROMPT = u('[PEXPECT_PROMPT+')
+PEXPECT_PROMPT = u'[PEXPECT_PROMPT>'
+PEXPECT_CONTINUATION_PROMPT = u'[PEXPECT_PROMPT+'
 
 class REPLWrapper(object):
     """Wrapper for a REPL.
@@ -39,7 +35,7 @@ class REPLWrapper(object):
                  continuation_prompt=PEXPECT_CONTINUATION_PROMPT,
                  extra_init_cmd=None):
         if isinstance(cmd_or_spawn, basestring):
-            self.child = pexpect.spawnu(cmd_or_spawn, echo=False)
+            self.child = pexpect.spawn(cmd_or_spawn, echo=False, encoding='utf-8')
         else:
             self.child = cmd_or_spawn
         if self.child.echo:
@@ -88,9 +84,11 @@ class REPLWrapper(object):
         if not cmdlines:
             raise ValueError("No command was given")
 
+        res = []
         self.child.sendline(cmdlines[0])
         for line in cmdlines[1:]:
-            self._expect_prompt(timeout=1)
+            self._expect_prompt(timeout=timeout)
+            res.append(self.child.before)
             self.child.sendline(line)
 
         # Command was fully submitted, now wait for the next prompt
@@ -100,15 +98,16 @@ class REPLWrapper(object):
             self._expect_prompt(timeout=1)
             raise ValueError("Continuation prompt found - input was incomplete:\n"
                              + command)
-        return self.child.before
+        return u''.join(res + [self.child.before])
 
 def python(command="python"):
     """Start a Python shell and return a :class:`REPLWrapper` object."""
-    return REPLWrapper(command, u(">>> "), u("import sys; sys.ps1={0!r}; sys.ps2={1!r}"))
+    return REPLWrapper(command, u">>> ", u"import sys; sys.ps1={0!r}; sys.ps2={1!r}")
 
 def bash(command="bash"):
     """Start a bash shell and return a :class:`REPLWrapper` object."""
     bashrc = os.path.join(os.path.dirname(__file__), 'bashrc.sh')
-    child = pexpect.spawnu(command, ['--rcfile', bashrc], echo=False)
-    return REPLWrapper(child, u'\$', u("PS1='{0}' PS2='{1}' PROMPT_COMMAND=''"),
+    child = pexpect.spawn(command, ['--rcfile', bashrc], echo=False,
+                          encoding='utf-8')
+    return REPLWrapper(child, u'\$', u"PS1='{0}' PS2='{1}' PROMPT_COMMAND=''",
                        extra_init_cmd="export PAGER=cat")
diff --git a/tests/test_replwrap.py b/tests/test_replwrap.py
index 28c7599..26ff53c 100644
--- a/tests/test_replwrap.py
+++ b/tests/test_replwrap.py
@@ -6,6 +6,8 @@ import os
 import pexpect
 from pexpect import replwrap
 
+skip_pypy = "This test fails on PyPy because of REPL differences"
+
 
 class REPLWrapTestCase(unittest.TestCase):
     def setUp(self):
@@ -25,10 +27,30 @@ class REPLWrapTestCase(unittest.TestCase):
         res = bash.run_command("time")
         assert 'real' in res, res
 
-        # PAGER should be set to cat, otherwise man hangs
+    def test_pager_as_cat(self):
+        " PAGER is set to cat, to prevent timeout in ``man sleep``. "
+        bash = replwrap.bash()
         res = bash.run_command('man sleep', timeout=5)
         assert 'SLEEP' in res, res
 
+    def test_long_running_multiline(self):
+        " ensure the default timeout is used for multi-line commands. "
+        bash = replwrap.bash()
+        res = bash.run_command("echo begin\r\nsleep 2\r\necho done")
+        self.assertEqual(res.strip().splitlines(), ['begin', 'done'])
+
+    def test_long_running_continuation(self):
+        " also ensure timeout when used within continuation prompts. "
+        bash = replwrap.bash()
+        # The two extra '\\' in the following expression force a continuation
+        # prompt:
+        # $ echo begin\
+        #     + ;
+        # $ sleep 2
+        # $ echo done
+        res = bash.run_command("echo begin\\\n;sleep 2\r\necho done")
+        self.assertEqual(res.strip().splitlines(), ['begin', 'done'])
+
     def test_multiline(self):
         bash = replwrap.bash()
         res = bash.run_command("echo '1 2\n3 4'")
@@ -47,7 +69,7 @@ class REPLWrapTestCase(unittest.TestCase):
         self.assertEqual(res.strip().splitlines(), ['1 2', '3 4'])
 
     def test_existing_spawn(self):
-        child = pexpect.spawnu("bash", timeout=5, echo=False)
+        child = pexpect.spawn("bash", timeout=5, echo=False, encoding='utf-8')
         repl = replwrap.REPLWrapper(child, re.compile('[$#]'),
                                     "PS1='{0}' PS2='{1}' "
                                     "PROMPT_COMMAND=''")
@@ -57,7 +79,7 @@ class REPLWrapTestCase(unittest.TestCase):
 
     def test_python(self):
         if platform.python_implementation() == 'PyPy':
-            raise unittest.SkipTest("This test fails on PyPy because of REPL differences")
+            raise unittest.SkipTest(skip_pypy)
 
         p = replwrap.python()
         res = p.run_command('4+7')
@@ -68,17 +90,16 @@ class REPLWrapTestCase(unittest.TestCase):
 
     def test_no_change_prompt(self):
         if platform.python_implementation() == 'PyPy':
-            raise unittest.SkipTest("This test fails on PyPy because of REPL differences")
+            raise unittest.SkipTest(skip_pypy)
 
-        child = pexpect.spawnu('python', echo=False, timeout=5)
+        child = pexpect.spawn('python', echo=False, timeout=5, encoding='utf-8')
         # prompt_change=None should mean no prompt change
-        py = replwrap.REPLWrapper(child, replwrap.u(">>> "), prompt_change=None,
-                                  continuation_prompt=replwrap.u("... "))
+        py = replwrap.REPLWrapper(child, u">>> ", prompt_change=None,
+                                  continuation_prompt=u"... ")
         assert py.prompt == ">>> "
 
         res = py.run_command("for a in range(3): print(a)\n")
         assert res.strip().splitlines() == ['0', '1', '2']
 
-
 if __name__ == '__main__':
     unittest.main()