Python Encapsulate a function to Print to a variable

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, inste开发者_运维知识库ad 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.

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!).

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

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()

Replace sys.stdout with a file-like object.

Use cStringIO ( see doc ).

from cStringIO import StringIO

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

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

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.