Add OS-level output capture, using file-descriptor redirection.

IPython/kernel/core/fd_redirector.py

@@ -0,0 +1,67 b''
+# encoding: utf-8
+
+"""
+Stdout/stderr redirector, at the OS level, using file descriptors.
+"""
+
+__docformat__ = "restructuredtext en"
+
+#-------------------------------------------------------------------------------
+#  Copyright (C) 2008  The IPython Development Team
+#
+#  Distributed under the terms of the BSD License.  The full license is in
+#  the file COPYING, distributed as part of this software.
+#-------------------------------------------------------------------------------
+
+
+import os
+import sys
+
+STDOUT = 1
+STDERR = 2
+
+class FDRedirector(object):
+    """ Class to redirect output (stdout or stderr) at the OS level using
+        file descriptors.
+    """ 
+
+    def __init__(self, fd=STDOUT):
+        """ fd is the file descriptor of the outpout you want to capture.
+            It can be STDOUT or STERR.
+        """
+        self.fd = fd
+        self.started = False
+
+    def start(self):
+        if not self.started:
+            self.oldhandle = os.dup(self.fd)
+            self.piper, self.pipew = os.pipe()
+            os.dup2(self.pipew, self.fd)
+            os.close(self.pipew)
+
+            self.started = True
+
+    def flush(self):
+        if self.fd == STDOUT:
+            sys.stdout.flush()
+        elif self.fd == STDERR:
+            sys.stderr.flush()
+
+    def stop(self):
+        if self.started:
+            self.flush()
+            os.dup2(self.oldhandle, self.fd)
+            os.close(self.oldhandle)
+            f = os.fdopen(self.piper, 'r')
+            output = f.read()
+            f.close()
+
+            self.started = False
+            return output
+        else:
+            return ''
+
+    def getvalue(self):
+        output = self.stop()
+        self.start()
+        return output

IPython/kernel/core/redirector_output_trap.py

@@ -0,0 +1,79 b''
+# encoding: utf-8
+
+""" 
+Trap stdout/stderr, including at the OS level. Calls a callback with
+the output each time Python tries to write to the stdout or stderr.
+"""
+
+__docformat__ = "restructuredtext en"
+
+#-------------------------------------------------------------------------------
+#  Copyright (C) 2008  The IPython Development Team
+#
+#  Distributed under the terms of the BSD License.  The full license is in
+#  the file COPYING, distributed as part of this software.
+#-------------------------------------------------------------------------------
+
+#-------------------------------------------------------------------------------
+# Imports
+#-------------------------------------------------------------------------------
+
+from fd_redirector import FDRedirector, STDOUT, STDERR
+
+from IPython.kernel.core.file_like import FileLike
+from IPython.kernel.core.output_trap import OutputTrap
+
+class RedirectorOutputTrap(OutputTrap):
+    """ Object which can trap text sent to stdout and stderr.
+    """
+
+    #------------------------------------------------------------------------
+    # OutputTrap interface.
+    #------------------------------------------------------------------------
+    def __init__(self, out_callback, err_callback):
+        # Callback invoked on write to stdout and stderr 
+        self.out_callback = out_callback
+        self.err_callback = err_callback
+
+        # File descriptor redirectors, to capture non-Python
+        # output.
+        self.out_redirector = FDRedirector(STDOUT)
+        self.err_redirector = FDRedirector(STDERR)
+
+        # Call the base class with file like objects that will trigger
+        # our callbacks
+        OutputTrap.__init__(self, out=FileLike(self.on_out_write),
+                                  err=FileLike(self.on_err_write), )
+
+
+    def set(self):
+        """ Set the hooks: set the redirectors and call the base class.
+        """
+        self.out_redirector.start()
+        #self.err_redirector.start()
+        OutputTrap.set(self)
+
+
+    def unset(self):
+        """ Remove the hooks: call the base class and stop the
+            redirectors.
+        """
+        OutputTrap.unset(self)
+        self.err_redirector.stop()
+        self.out_redirector.stop()
+
+
+    #------------------------------------------------------------------------
+    # Callbacks for synchronous output 
+    #------------------------------------------------------------------------
+    def on_out_write(self, string):
+        """ Callback called when there is some Python output on stdout.
+        """
+        self.out_callback(self.out_redirector.getvalue() + string)
+
+    def on_err_write(self, string):
+        """ Callback called when there is some Python output on stderr.
+        """
+        self.err_callback(self.err_redirector.getvalue() + string)
+
+

IPython/kernel/core/tests/test_redirectors.py

@@ -0,0 +1,50 b''
+"""
+Test the output capture at the OS level, using file descriptors.
+"""
+
+import os
+from cStringIO import StringIO
+
+
+def test_redirector():
+    """ Checks that the redirector can be used to do synchronous capture.
+    """
+    from IPython.kernel.core.fd_redirector import FDRedirector
+    r = FDRedirector()
+    out = StringIO()
+    try:
+        r.start()
+        for i in range(10):
+            os.system('echo %ic' % i)
+            print >>out, r.getvalue(),
+            print >>out, i
+    except:
+        r.stop()
+        raise
+    r.stop()
+    assert out.getvalue() == "".join("%ic\n%i\n" %(i, i) for i in range(10))
+
+
+def test_redirector_output_trap():
+    """ This test check not only that the redirector_output_trap does
+        trap the output, but also that it does it in a gready way, that
+        is by calling the callabck ASAP.
+    """
+    from IPython.kernel.core.redirector_output_trap import RedirectorOutputTrap
+    out = StringIO()
+    trap = RedirectorOutputTrap(out.write, out.write)
+    try:
+        trap.set()
+        for i in range(10):
+            os.system('echo %ic' % i)
+            print "%ip" % i
+            print >>out, i
+    except:
+        trap.unset()
+        raise
+    trap.unset()
+    assert out.getvalue() == "".join("%ic\n%ip\n%i\n" %(i, i, i) 
+                                                    for i in range(10))
+
+    
+