Отключить все распечатки, кроме сделанных в вызываемой функции

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)

Но это включает в себя изменение кода функции и кажется немного хакерским.В идеале я хотел бы сделать это, используя декоратор, без необходимости изменять код функции.

Наиболее естественным способом было бы найти отпечатки, не вызванные в my_func, и заставить их выводить на os.devnull вобертка.Но я не могу найти, как это сделать.Заранее спасибо.

Answer 1

@ blhsing дал хороший подход.Но я решил проблему, используя pprint.pprint в своей функции и отключив выходы print.Это работает, потому что pprint использует более низкий уровень stream.write("...") вместо начальной загрузки на print.Код:

def disable_prints(f):
    @functools.wraps(f)
    def decorated(*args, **kwargs):
        global print
        # Disable all calls made to print(...) by relacing stdout with devnull
        print = functools.partial(print, file=open(os.devnull, 'w'))
        f_returns = f(*args, **kwargs)
        # Restore back
        print = functools.partial(print, file=sys.stdout)
        return f_returns
    return decorated

, чтобы:

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

@disable_prints
def my_func():
    pprint("Starting inside my_func")
    func()
    pprint("In my_func")
    func()

my_func()

выводил:

Starting inside my_func 
In my_func

Answer 2

Вы можете сохранить ссылку на функцию 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