Python：爲特定函數調用修補打印函數？ （用於打印遞歸樹的裝飾器）

Question

我寫了一個裝飾器來打印由某個函數調用產生的遞歸樹。

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 

實例：

@printRecursionTree 
def fib(n): 
    if n <= 1: 
     print('Base Case') 
     return n 
    print('Recursive Case') 
    return fib(n-1) + fib(n-2) 

# This works with mutually recursive functions too, 
# since the variable _recursiondepth is global 
@printRecursionTree 
def iseven(n): 
    print('checking if even') 
    if n == 0: return True 
    return isodd(n-1) 

@printRecursionTree 
def isodd(n): 
    print('checking if odd') 
    if n == 0: return False 
    return iseven(n-1) 

iseven(5) 
fib(5) 

'''Prints: 

└── iseven(5): 
    │▒▒ checking if even 
    │▒▒ Note how the print 
    │▒▒ statements get nicely indented 
    ├── isodd(4): 
    │ │▒▒ checking if odd 
    │ ├── iseven(3): 
    │ │ │▒▒ checking if even 
    │ │ │▒▒ Note how the print 
    │ │ │▒▒ statements get nicely indented 
    │ │ ├── isodd(2): 
    │ │ │ │▒▒ checking if odd 
    │ │ │ ├── iseven(1): 
    │ │ │ │ │▒▒ checking if even 
    │ │ │ │ │▒▒ Note how the print 
    │ │ │ │ │▒▒ statements get nicely indented 
    │ │ │ │ ├── isodd(0): 
    │ │ │ │ │ │▒▒ checking if odd 
    │ │ │ │ │ ╰ False 
    │ │ │ │ ╰ False 
    │ │ │ ╰ False 
    │ │ ╰ False 
    │ ╰ False 
    ╰ False 

└── fib(5): 
    │▒▒ Recursive Case 
    ├── fib(4): 
    │ │▒▒ Recursive Case 
    │ ├── fib(3): 
    │ │ │▒▒ Recursive Case 
    │ │ ├── fib(2): 
    │ │ │ │▒▒ Recursive Case 
    │ │ │ ├── fib(1): 
    │ │ │ │ │▒▒ Base Case 
    │ │ │ │ ╰ 1 
    │ │ │ ├── fib(0): 
    │ │ │ │ │▒▒ Base Case 
    │ │ │ │ ╰ 0 
    │ │ │ ╰ 1 
    │ │ ├── fib(1): 
    │ │ │ │▒▒ Base Case 
    │ │ │ ╰ 1 
    │ │ ╰ 2 
    │ ├── fib(2): 
    │ │ │▒▒ Recursive Case 
    │ │ ├── fib(1): 
    │ │ │ │▒▒ Base Case 
    │ │ │ ╰ 1 
    │ │ ├── fib(0): 
    │ │ │ │▒▒ Base Case 
    │ │ │ ╰ 0 
    │ │ ╰ 1 
    │ ╰ 3 
    ├── fib(3): 
    │ │▒▒ Recursive Case 
    │ ├── fib(2): 
    │ │ │▒▒ Recursive Case 
    │ │ ├── fib(1): 
    │ │ │ │▒▒ Base Case 
    │ │ │ ╰ 1 
    │ │ ├── fib(0): 
    │ │ │ │▒▒ Base Case 
    │ │ │ ╰ 0 
    │ │ ╰ 1 
    │ ├── fib(1): 
    │ │ │▒▒ Base Case 
    │ │ ╰ 1 
    │ ╰ 2 
    ╰ 5 
''' 

此示例代碼工作正常，只要它是在將裝飾定義相同的文件。

但是，如果從某個模塊導入裝飾器，則打印語句不再縮進。

我知道這種行爲的產生是因爲由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