Python：为特定的函数调用修补print函数？（用于打印递归树的装饰器）

Question

我编写了一个decorator来打印由某个函数调用生成的递归树。在

from functools import wraps

def printRecursionTree(func):
    global _recursiondepth
    _print = print
    _recursiondepth = 0

    def getpads():
        if _recursiondepth == 0:
            strFn    = '{} └──'.format(' │  ' * (_recursiondepth-1))
            strOther = '{}  ▒▒'.format(' │  ' * (_recursiondepth-1))
            strRet   = '{}    '.format(' │  ' * (_recursiondepth-1))
        else:
            strFn    = '    {} ├──'.format(' │  ' * (_recursiondepth-1))
            strOther = '    {} │▒▒'.format(' │  ' * (_recursiondepth-1))
            strRet   = '    {} │  '.format(' │  ' * (_recursiondepth-1))

        return strFn, strRet, strOther

    def indentedprint():
        @wraps(print)
        def wrapper(*args, **kwargs):
            strFn, strRet, strOther = getpads()
            _print(strOther, end=' ')
            _print(*args, **kwargs)
        return wrapper


    @wraps(func)
    def wrapper(*args, **kwargs):
        global _recursiondepth
        global print

        strFn, strRet, strOther = getpads()

        if args and kwargs:
            _print(strFn, '{}({}, {}):'.format(func.__qualname__, ', '.join(args), kwargs))
        else:
            _print(strFn, '{}({}):'.format(func.__qualname__, ', '.join(map(str, args)) if args else '', kwargs if kwargs else ''))
        _recursiondepth += 1
        print, backup = indentedprint(), print
        retval = func(*args, **kwargs)
        print = backup
        _recursiondepth -= 1
        _print(strRet, '╰', retval)
        if _recursiondepth == 0:
            _print()
        return retval

    return wrapper

用法示例：

^{pr2}$

这个示例代码只要位于定义decorator的同一个文件中，就可以正常工作。在

但是，如果从某个模块导入decorator，print语句将不再缩进。在

我理解这种行为是因为decorator修补的print语句是它自己模块的全局语句，而不是在模块之间共享的。在

1. 我怎么解决这个问题？在
2. 有没有更好的方法只为对另一个函数的特定调用修补一个函数？在

Answer 1

您可以通过在builtins模块中替换它来更改所有模块的内置打印函数的行为。在

因此，将赋值改为全局变量print，并将赋值改为builtins.print（在导入builtins之后）：

import builtins

...

    @wraps(func)
    def wrapper(*args, **kwargs):
        global _recursiondepth # no more need for global print up here

        strFn, strRet, strOther = getpads()

        if args and kwargs:
            _print(strFn, '{}({}, {}):'.format(func.__qualname__, ', '.join(args), kwargs))
        else:
            _print(strFn, '{}({}):'.format(func.__qualname__, ', '.join(map(str, args)) if args else '', kwargs if kwargs else ''))
        _recursiondepth += 1
        builtins.print, backup = indentedprint(), print   # change here
        retval = func(*args, **kwargs)
        builtins.print = backup                           # and here
        _recursiondepth -= 1
        _print(strRet, '╰', retval)
        if _recursiondepth == 0:
            _print()
        return retval