Python Encapsulate a function to Print to a variable

Question

If I have a function that contains a lot of print statements:

ie.

def funA():
  print "Hi"
  print "There"
  print "Friend"
  print "!"

What I want to do is something like this

def main():
  ##funA() does not print to screen here
  a = getPrint(funA()) ##where getPrint is some made up function/object
  print a ##prints what funA would normally print at this step

So when funcA gets called it doesn't do any printing, instead it output to an object. I then print the object to get the result. Is there a way of doing this? I also do not want to touch the original function.

I hope it makes sense.

related: http://stackoverflow.com/questions/1218933/can-i-redirect-the-stdout-in-python-into-some-sort-of-string-buffer — jldupont, Dec 31 '09 at 00:59

Alex Martelli · Accepted Answer · 2010-10-04T04:39:17.553

8

You can do almost exactly what you want, as long as you don't mind a tiny syntax difference:

import cStringIO
import sys

def getPrint(thefun, *a, **k):
  savstdout = sys.stdout
  sys.stdout = cStringIO.StringIO()
  try:
    thefun(*a, **k)
  finally:
    v = sys.stdout.getvalue()
    sys.stdout = savstdout
  return v

The tiny difference is that you must call getPrint(funA), not getPrint(funA()) -- i.e., you must pass the function object itself, without the trailing parentheses that would call it immediately, before getPrint can do its magic.

If you absolutely insist on those extra parentheses, then getPrint cannot do all the needed preparation, and must be supplemented by other code to prepare things right (I strongly recommend losing the extra parentheses, thus enabling the encapsulation of all the functionality inside getPrint!).

edited Oct 04 '10 at 04:39

answered Dec 31 '09 at 01:10

Alex Martelli

854,459
170
1,222
1,395

+1. I would use sys.__stdout__, which officially contains the original sys.stdout. There is no need to save the original sys.stdout in savstdout, if I'm not mistaken. – Eric O. Lebigot Dec 31 '09 at 08:37
@EOL, if you're running your code under IDLE or other GUI-based IDE, "the original stdout" may be useless or unusable (e.g. it may be closed, directed to a non-existing terminal, etc); and there IS most definitely a need to save and restore `sys.stdout` in order to put the IDE back to a usable state when `getPrint` is gone. – Alex Martelli Dec 31 '09 at 17:10
1

Why are you calling the cStringIO module? – Oct 04 '10 at 04:35

score 3 · Answer 2 · answered Dec 31 '09 at 01:12

from cStringIO import StringIO

def getPrint(func, *args, **kwds):
  old_stdout = sys.stdout
  sys.stdout = StringIO()
  try:
    func(*args, **kwds)
  except:
    raise
  else:
    return sys.stdout.getvalue()
  finally:
    sys.stdout = old_stdout

#...
a = getPrint(funA) # notice no (), it is called by getPrint
print a.rstrip("\n") # avoid extra trailing lines

score 2 · Answer 3 · answered Dec 31 '09 at 01:56

Best way is to do a context manager

from contextlib import contextmanager
import StringIO
import sys

@contextmanager
def capture():
    old_stdout = sys.stdout
    sys.stdout = StringIO.StringIO()
    try:
        yield sys.stdout
    finally:
        sys.stdout = old_stdout

Now you can run any printing code:

with capture() as c:
    funA()
    funB()
    print 'HELLO!'

then later:

print c.getvalue()

score 1 · Answer 4 · answered Dec 31 '09 at 00:58

1

Replace sys.stdout with a file-like object.

answered Dec 31 '09 at 00:58

Ignacio Vazquez-Abrams

776,304
153
1,341
1,358

score 1 · Answer 5 · answered Dec 31 '09 at 01:00

1

Use cStringIO ( see doc ).

from cStringIO import StringIO

old_stdout = sys.stdout
sys.stdout = mystdout = StringIO()

getPrint( funA() )
# use mystdout to get string

answered Dec 31 '09 at 01:00

jldupont

93,734
56
203
318

score 0 · Answer 6 · answered Dec 31 '09 at 02:19

0

The simplest thing is to change your funA() to not print anything, but simply to return the string values.

Like so:

def funA():
    return "Hi\n" + "There\n" + "Friend\n" + "!\n"

# later:
print(funA())

It's always easy to collect strings and print them; it's tricker to to collect strings as they are being printed.

If you have a huge body of existing printing functions, then yeah, use one of the tricks provided here to collect the output.

answered Dec 31 '09 at 02:19

steveha

74,789
21
92
117

the problem lies in the fact that the function is not trivial as the one I posted. It calls upon other functions that calls upon other functions that does the printing. – John Jiang Jan 03 '10 at 13:17
Well, you have two choices. If all the functions use a common single function to emit their output, then hook that one function and you can capture all the output; otherwise (maybe they all use the Python 2.x print *statement* which you can't really hook) then you will want to use one of the tricks documented on this page. – steveha Jan 03 '10 at 21:59

Python Encapsulate a function to Print to a variable

6 Answers6