如何使一个函数的输出静音到控制台？

Question

我有几个函数，我只需要最外层函数调用的输出。例如：

def f1():
    print 'foo'

def f2(boo=True):
    f1()
    if boo:
        print 'bar'     
    else:
        print 'black sheep'

def f3():
    f2()
    print 'shh!!!'  

f2(True)
print
f2(False)
print 
f3()

以上脚本输出：

^{pr2}$

所需输出：

bar

black sheep

shh!!!

Answer 1

使用inspect

import inspect
def f1():
    if (len(inspect.stack())) <= 2:
        print 'foo'

def f2(boo=True):
    f1()
    if boo:
        if (len(inspect.stack())) <= 2:
            print 'bar'     
    else:
        if (len(inspect.stack())) <= 2:
            print 'black sheep'

def f3():
    f2()
    if (len(inspect.stack())) <= 2:
        print 'shh!!!'  

f2(True)
print
f2(False)
print 
f3()

Answer 2

如果允许修改f2，f3内的代码，那么一种方法是使用自定义上下文管理器，它将sys.stdout分配给其他内容，并在上下文管理器退出时将sys.stdout重新分配给原始STDOUT。在

import sys, StringIO

class Supress_print(object):

    def __init__(self):
        pass

    def __enter__(self):
        self.stdout = sys.stdout
        sys.stdout = StringIO.StringIO()

    def __exit__(self, *args):
        sys.stdout = self.stdout

def f1():
    print 'foo'

def f2(boo=True):
    with Supress_print():
        f1()
    if boo:
        print 'bar'
    else:
        print 'black sheep'

def f3():
    with Supress_print():
        f2()
    print 'shh!!!'  

f2(True)
print
f2(False)
print 
f3()

输出：

^{pr2}$

---

更新：

import sys, inspect, functools, StringIO

def supress_print(func):

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        #print inspect.getouterframes(inspect.currentframe())[3][3], func.__name__
        if inspect.getouterframes(inspect.currentframe())[3][3] != 'main':
            stdout = sys.stdout
            sys.stdout = StringIO.StringIO()
            val = func(*args, **kwargs)
            sys.stdout = stdout
            return val
        else:
            return func(*args, **kwargs)

    return wrapper

@supress_print
def f1():
    print 'foo'


@supress_print
def f2(boo=True):
    f1()
    if boo:
        print 'bar'
    else:
        print 'black sheep'

def f3():
    f2()
    print 'shh!!!'