From 97262e9521f384fe8f55566bdfa1b4794e55bd8b 2008-08-01 05:04:44 From: Gael Varoquaux Date: 2008-08-01 05:04:44 Subject: [PATCH] Add OS-level output capture, using file-descriptor redirection. --- diff --git a/IPython/kernel/core/fd_redirector.py b/IPython/kernel/core/fd_redirector.py new file mode 100644 index 0000000..0941fac --- /dev/null +++ b/IPython/kernel/core/fd_redirector.py @@ -0,0 +1,67 @@ +# 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 diff --git a/IPython/kernel/core/redirector_output_trap.py b/IPython/kernel/core/redirector_output_trap.py new file mode 100644 index 0000000..23fe145 --- /dev/null +++ b/IPython/kernel/core/redirector_output_trap.py @@ -0,0 +1,79 @@ +# 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) + + diff --git a/IPython/kernel/core/tests/test_redirectors.py b/IPython/kernel/core/tests/test_redirectors.py new file mode 100644 index 0000000..85310bd --- /dev/null +++ b/IPython/kernel/core/tests/test_redirectors.py @@ -0,0 +1,50 @@ +""" +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)) + + +