python - 禁用所有打印，除了在调用函数中完成的打印

Question

我有一个函数，我从同一个模块或其他一些模块中调用一些函数：

from __future__ import print_function

def func():
    print("Inside func")

def my_func():
    print("Starting inside my_func ")
    func()
    print("In my_func")
    func()

执行my_func输出：

Starting inside my_func
Inside func
In my_func
Inside func

但我只想看看

Starting inside my_func
In my_func

所以，我想禁用所有打印，除了直接在函数中制作的打印my_func()。这可能包括对函数的递归调用。所以用堆栈级别做一些事情是行不通的。

我可以考虑这样做

def func():
    print("Inside func")

def my_print(*args):
    print(*args)

def my_func():
    global my_print, print
    my_print("Starting inside my_func ")
    print = functools.partial(print, file=open(os.devnull, 'w'))
    func()
    print = functools.partial(print, file=sys.stdout)
    my_print("In my_func")
    print = functools.partial(print, file=open(os.devnull, 'w'))
    func()
    print = functools.partial(print, file=sys.stdout)

但这涉及修改功能代码，似乎有点hacky。理想情况下，我希望使用装饰器来执行此操作，而无需修改功能代码。

最自然的方法是找到未调用的打印my_func并将它们输出到os.devnull包装器中。但我找不到如何做到这一点。提前致谢。

标签： pythonprintingdecorator

Answer 1

您可以将对该print函数的引用保存在一个变量中orig_print并用一个什么都不做的函数覆盖print，然后在您希望允许打印的函数上使用一个装饰器，print以orig_print使用ast.NodeTransformer子类重命名所有调用：

from __future__ import print_function
import inspect
import ast
from textwrap import dedent

orig_print = print
print = lambda *args, **kwargs: None

class EnablePrint(ast.NodeTransformer):
    # remove the enable_print decorator from the decorator list so the transformed
    # function won't be re-decorated when executed
    def visit_FunctionDef(self, node):
        node.decorator_list = [
            decorator for decorator in node.decorator_list
            if not isinstance(decorator, ast.Name) or decorator.id != 'enable_print'
        ]
        self.generic_visit(node)
        return node

    def visit_Call(self, node):
        if node.func.id == 'print':
            node.func.id = 'orig_print'
        return node

def enable_print(func):
    node = ast.parse(dedent(inspect.getsource(func)))
    EnablePrint().visit(node)
    scope = {}
    exec(compile(node, inspect.getfile(func), 'exec'), func.__globals__, scope)
    return scope[func.__name__]

以便：

def func():
    print("Inside func")

@enable_print
def my_func():
    print("Starting inside my_func ")
    func()
    print("In my_func")
    func()

my_func()

会输出：

Starting inside my_func 
In my_func