|
|
|
# -*- python -*-
|
|
|
|
#
|
|
|
|
# wpylib.py.stdout_capture
|
|
|
|
# Created: 20110809
|
|
|
|
# Wirawan Purwanto
|
|
|
|
#
|
|
|
|
# Hack tools for capturing sys.stdout in a function call.
|
|
|
|
#
|
|
|
|
"""
|
|
|
|
wpylib.py.stdout_capture
|
|
|
|
|
|
|
|
Capture tool for sys.stdout output, for interactive uses (primarily,
|
|
|
|
so that the output from sys.stdout is recorded to an ipython log file).
|
|
|
|
|
|
|
|
Example uses:
|
|
|
|
|
|
|
|
1) redirecting stdout explicitly
|
|
|
|
|
|
|
|
from wpylib.py.stdout_capture import StdoutCapture
|
|
|
|
xstdout = StdoutCapture()
|
|
|
|
# the original command was: fc.calculate_and_print_energy(h, UHF_01, UHF_01, out=sys.stdout)
|
|
|
|
# ^^^^^^^^^^^^^^
|
|
|
|
# this now becomes:
|
|
|
|
xstdout(fc.calculate_and_print_energy, h, UHF_01, UHF_01, out=xstdout.stdout)
|
|
|
|
|
|
|
|
|
|
|
|
Original-source: `ipython_common.py`.
|
|
|
|
"""
|
|
|
|
|
|
|
|
import sys
|
|
|
|
|
|
|
|
class StdoutCapture(object):
|
|
|
|
"""Capture tool for sys.stdout output, for interactive uses (primarily,
|
|
|
|
so that the output from sys.stdout is recorded to an ipython log file).
|
|
|
|
|
|
|
|
"""
|
|
|
|
class outfile_proxy(object):
|
|
|
|
"""A thin proxy for output file.
|
|
|
|
The actual object, self.obj, *MUST* be set below."""
|
|
|
|
def __init__(self):
|
|
|
|
pass
|
|
|
|
def write(self, stuff):
|
|
|
|
self.obj.write(stuff)
|
|
|
|
def flush(self):
|
|
|
|
self.obj.flush()
|
|
|
|
def close(self):
|
|
|
|
pass
|
|
|
|
|
|
|
|
def __init__(self):
|
|
|
|
pass
|
|
|
|
|
|
|
|
@property
|
|
|
|
def stdout(self):
|
|
|
|
"""A stdout proxy that works with StdoutCapture."""
|
|
|
|
if not hasattr(self, "my_stdout_proxy"):
|
|
|
|
self.my_stdout_proxy = self.outfile_proxy()
|
|
|
|
self.my_stdout_proxy_created = 1
|
|
|
|
return self.my_stdout_proxy
|
|
|
|
|
|
|
|
def Capture(self, proc, *args, **argkw):
|
|
|
|
'''Captures sys.stdout while running proc (and its child subroutines).
|
|
|
|
Returns a `printstr` object which, if displayed "as is" on an
|
|
|
|
interactive python shell, would be logged in the log file as well.
|
|
|
|
'''
|
|
|
|
from StringIO import StringIO
|
|
|
|
from wpylib.interactive_tools import printstr
|
|
|
|
|
|
|
|
needs_stdout_redir = not hasattr(StdoutCapture, "_StdoutCapture_save_stdout")
|
|
|
|
my_stdout_proxy_created = getattr(self, "my_stdout_proxy_created", 0)
|
|
|
|
|
|
|
|
if needs_stdout_redir:
|
|
|
|
#print "stdout not redirected yet"
|
|
|
|
sys.stdout.flush()
|
|
|
|
temp_stdout = StringIO()
|
|
|
|
StdoutCapture._StdoutCapture_save_stdout = sys.stdout
|
|
|
|
sys.stdout = temp_stdout
|
|
|
|
|
|
|
|
if my_stdout_proxy_created:
|
|
|
|
# delete it so it won't be touched later
|
|
|
|
delattr(self, "my_stdout_proxy_created")
|
|
|
|
if hasattr(self, "my_stdout_proxy"):
|
|
|
|
self.my_stdout_proxy.obj = temp_stdout
|
|
|
|
|
|
|
|
rslt1 = None # print output
|
|
|
|
rslt2 = None # function output
|
|
|
|
try:
|
|
|
|
#print >> sys.stderr, "calling proc"
|
|
|
|
rslt2 = proc(*args, **argkw)
|
|
|
|
#print >> sys.stderr, "end proc"
|
|
|
|
except:
|
|
|
|
#print >> sys.stderr, "finally proc"
|
|
|
|
if needs_stdout_redir:
|
|
|
|
sys.stdout = StdoutCapture._StdoutCapture_save_stdout
|
|
|
|
delattr(StdoutCapture, "_StdoutCapture_save_stdout")
|
|
|
|
rslt1 = temp_stdout.getvalue()
|
|
|
|
temp_stdout.close()
|
|
|
|
|
|
|
|
# clean up the stdout proxy:
|
|
|
|
if my_stdout_proxy_created:
|
|
|
|
# delete it so it won't be touched later
|
|
|
|
delattr(self, "my_stdout_proxy")
|
|
|
|
if rslt1:
|
|
|
|
print >> sys.stderr, rslt1
|
|
|
|
print >> sys.stderr, "Exception caught while capturing stdout"
|
|
|
|
raise
|
|
|
|
|
|
|
|
else:
|
|
|
|
#print >> sys.stderr, "finally proc"
|
|
|
|
if needs_stdout_redir:
|
|
|
|
sys.stdout = StdoutCapture._StdoutCapture_save_stdout
|
|
|
|
delattr(StdoutCapture, "_StdoutCapture_save_stdout")
|
|
|
|
print >> sys.stderr, temp_stdout.tell()
|
|
|
|
rslt1 = temp_stdout.getvalue()
|
|
|
|
#print >> sys.stderr, "stdout redirected back"
|
|
|
|
#sys.stdout.flush()
|
|
|
|
temp_stdout.close()
|
|
|
|
|
|
|
|
# clean up the stdout proxy:
|
|
|
|
if my_stdout_proxy_created:
|
|
|
|
# delete it so it won't be touched later
|
|
|
|
delattr(self, "my_stdout_proxy")
|
|
|
|
|
|
|
|
if rslt2 != None:
|
|
|
|
rslt = printstr("".join([rslt1, "\n", repr(rslt2)]))
|
|
|
|
rslt.rslt = rslt2 # for user to fetch later
|
|
|
|
else:
|
|
|
|
rslt = printstr(rslt1)
|
|
|
|
return rslt
|
|
|
|
else:
|
|
|
|
return rslt2
|
|
|
|
|
|
|
|
def __call__(self, proc, *args, **argkw):
|
|
|
|
return self.Capture(proc, *args, **argkw)
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
def test_Capture(arg):
|
|
|
|
"""Test routine for capture.
|
|
|
|
Example usage:
|
|
|
|
|
|
|
|
>>> cap = StdoutCapture()
|
|
|
|
>>> cap.Capture(cap.test_Capture, 5)"""
|
|
|
|
print "sys.stdout object is:", repr(sys.stdout)
|
|
|
|
a = 0
|
|
|
|
for i in range(arg):
|
|
|
|
print "i = ", i, " I want ", ("a" * i)
|
|
|
|
a+=i
|
|
|
|
return a
|