Merge pull request #51 from takluyver/replwrap

Add high level API for wrapping REPLs

5 files changed, 167 insertions, 3 deletions

diff --git a/.travis.yml b/.travis.yml
index 15ba2a2..601fb3a 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -7,8 +7,8 @@ python:
   - 3.3
   - pypy
 
-virtualenv:
-  system_site_packages: true
+#virtualenv:
+#  system_site_packages: true
 
 before_install:
   - sudo apt-get install python-yaml python3-yaml
diff --git a/doc/api/index.rst b/doc/api/index.rst
index b902131..1a6a6ae 100644
--- a/doc/api/index.rst
+++ b/doc/api/index.rst
@@ -6,6 +6,7 @@ API documentation
 
    pexpect
    fdpexpect
+   replwrap
    pxssh
    screen
-   ANSI
\ No newline at end of file
+   ANSI
diff --git a/doc/api/replwrap.rst b/doc/api/replwrap.rst
new file mode 100644
index 0000000..bf44a94
--- /dev/null
+++ b/doc/api/replwrap.rst
@@ -0,0 +1,26 @@
+replwrap - Control read-eval-print-loops
+========================================
+
+.. automodule:: pexpect.replwrap
+
+.. versionadded:: 3.3
+
+.. autoclass:: REPLWrapper
+
+   .. automethod:: run_command
+
+.. data:: PEXPECT_PROMPT
+
+   A string that can be used as a prompt, and is unlikely to be found in output.
+
+Using the objects above, it is easy to wrap a REPL. For instance, to use a
+Python shell::
+
+    py = REPLWrapper("python", ">>> ", "import sys; sys.ps1={!r}; sys.ps2={!r}")
+    py.run_command("4+7")
+
+Convenience functions are provided for Python and bash shells:
+
+.. autofunction:: python
+
+.. autofunction:: bash
diff --git a/pexpect/replwrap.py b/pexpect/replwrap.py
new file mode 100644
index 0000000..965c790
--- /dev/null
+++ b/pexpect/replwrap.py
@@ -0,0 +1,98 @@
+"""Generic wrapper for read-eval-print-loops, a.k.a. interactive shells
+"""
+import signal
+import sys
+
+import pexpect
+
+PY3 = (sys.version_info[0] >= 3)
+
+if PY3:
+    def u(s): return s
+else:
+    def u(s): return s.decode('utf-8')
+
+PEXPECT_PROMPT = u('[PEXPECT_PROMPT>')
+PEXPECT_CONTINUATION_PROMPT = u('[PEXPECT_PROMPT+')
+
+class REPLWrapper(object):
+    """Wrapper for a REPL.
+    
+    :param cmd_or_spawn: This can either be an instance of :class:`pexpect.spawn`
+      in which a REPL has already been started, or a str command to start a new
+      REPL process.
+    :param str orig_prompt: The prompt to expect at first.
+    :param str prompt_change: A command to change the prompt to something more
+      unique. If this is ``None``, the prompt will not be changed. This will
+      be formatted with the new and continuation prompts as positional
+      parameters, so you can use ``{}`` style formatting to insert them into
+      the command.
+    :param str new_prompt: The more unique prompt to expect after the change.
+    """
+    def __init__(self, cmd_or_spawn, orig_prompt, prompt_change,
+                 new_prompt=PEXPECT_PROMPT,
+                 continuation_prompt=PEXPECT_CONTINUATION_PROMPT):
+        if isinstance(cmd_or_spawn, str):
+            self.child = pexpect.spawnu(cmd_or_spawn)
+        else:
+            self.child = cmd_or_spawn
+        self.child.setecho(False)  # Don't repeat our input.
+        
+        if prompt_change is None:
+            self.prompt = orig_prompt
+        else:
+            self.set_prompt(orig_prompt,
+                        prompt_change.format(new_prompt, continuation_prompt))
+            self.prompt = new_prompt
+        self.continuation_prompt = continuation_prompt
+
+        self._expect_prompt()
+
+    def set_prompt(self, orig_prompt, prompt_change):
+        self.child.expect_exact(orig_prompt)
+        self.child.sendline(prompt_change)
+
+    def _expect_prompt(self, timeout=-1):
+        return self.child.expect_exact([self.prompt, self.continuation_prompt],
+                                       timeout=timeout)
+    
+    def run_command(self, command, timeout=-1):
+        """Send a command to the REPL, wait for and return output.
+        
+        :param str command: The command to send. Trailing newlines are not needed.
+          This should be a complete block of input that will trigger execution;
+          if a continuation prompt is found after sending input, :exc:`ValueError`
+          will be raised.
+        :param int timeout: How long to wait for the next prompt. -1 means the
+          default from the :class:`pexpect.spawn` object (default 30 seconds).
+          None means to wait indefinitely.
+        """
+        # Split up multiline commands and feed them in bit-by-bit
+        cmdlines = command.splitlines()
+        # splitlines ignores trailing newlines - add it back in manually
+        if command.endswith('\n'):
+            cmdlines.append('')
+        if not cmdlines:
+            raise ValueError("No command was given")
+
+        self.child.sendline(cmdlines[0])
+        for line in cmdlines[1:]:
+            self._expect_prompt(timeout=1)
+            self.child.sendline(line)
+
+        # Command was fully submitted, now wait for the next prompt
+        if self._expect_prompt(timeout=timeout) == 1:
+            # We got the continuation prompt - command was incomplete
+            self.child.kill(signal.SIGINT)
+            self._expect_prompt(timeout=1)
+            raise ValueError("Continuation prompt found - input was incomplete:\n"
+                             + command)
+        return 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}"))
+
+def bash(command="bash", orig_prompt=u("$")):
+    """Start a bash shell and return a :class:`REPLWrapper` object."""
+    return REPLWrapper(command, orig_prompt, u("PS1='{0}'; PS2='{1}'"))
diff --git a/tests/test_replwrap.py b/tests/test_replwrap.py
new file mode 100644
index 0000000..835bf63
--- /dev/null
+++ b/tests/test_replwrap.py
@@ -0,0 +1,39 @@
+import sys
+import unittest
+
+import pexpect
+from pexpect import replwrap
+
+class REPLWrapTestCase(unittest.TestCase):
+    def test_python(self):
+        bash = replwrap.bash()
+        res = bash.run_command("time")
+        assert 'real' in res, res
+
+    def test_multiline(self):
+        bash = replwrap.bash()
+        res = bash.run_command("echo '1 2\n3 4'")
+        self.assertEqual(res.strip().splitlines(), ['1 2', '3 4'])
+
+        # Should raise ValueError if input is incomplete
+        try:
+            bash.run_command("echo '5 6")
+        except ValueError:
+            pass
+        else:
+            assert False, "Didn't raise ValueError for incomplete input"
+
+        # Check that the REPL was reset (SIGINT) after the incomplete input
+        res = bash.run_command("echo '1 2\n3 4'")
+        self.assertEqual(res.strip().splitlines(), ['1 2', '3 4'])
+
+    def test_existing_spawn(self):
+        child = pexpect.spawnu("bash")
+        repl = replwrap.REPLWrapper(child, replwrap.u("$ "),
+                            "PS1='{0}'; PS2='{1}'")
+
+        res = repl.run_command("echo $HOME")
+        assert res.startswith('/'), res
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file