Book/파이썬 코딩의 기술

[BW 26] functools.wrap을 사용해 함수 데코레이터를 정의하라

31514 2024. 10. 25. 11:47

데코레이터는 자신이 감싸고 있는 함수가 호출되기 전과 후에 코드를 추가로 실행해준다.

즉, 데코레이터가 자신이 감싸고 있는 함수의 입력 인자, 반환 값, 함수에서 발생한 오류에 접근할 수 있다는 뜻이다.

 

예를 들어 함수가 호출될 때마다 인자 값과 반환 값을 출력하는 함수가 있다고 하자.

def trace(func):
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        print(f'{func.__name__}({args!r}, {kwargs!r}'
              f'-> {result!r}')
        return result
    return wrapper

 

그리고 피보나치 함수를 정의하고 trace 함수의 인자로 피보나치 함수를 입력할 수도 있지만,

def fibonacci(n):
    """n번째 피보나치 수를 반환한다."""
    if n in (0, 1):
        return n
    return (fibonacci(n - 2) + fibonacci(n - 1))
    
fibonacci = trace(fibonacci)

 

 

애초에 함수를 정의할 때, 데코레이터를 사용할 수도 있다.

@trace
def fibonacci(n):
    """n번째 피보나치 수를 반환한다."""
    if n in (0, 1):
        return n
    return (fibonacci(n - 2) + fibonacci(n - 1))

 

이를 실행하면 다음과 같다.

fibonacci(3)

>>>
fibonacci((1,), {}-> 1
fibonacci((0,), {}-> 0
fibonacci((1,), {}-> 1
fibonacci((2,), {}-> 1
fibonacci((3,), {}-> 2

 

하지만 데코레이터를 사용한 함수의 이름이 피보나치가 아닌게 된다.

print(fibonacci)

>>>
<function trace.<locals>.wrapper at 0x1049f0b80>

 



만약 데코레이터를 사용하지 않았더라면 다음과 같은 결과가 나왔을 것이다.

<function fibonacci at 0x10288ce00>

 

 

그리고 help 내장 함수의 경우도 마찬가지다. 피보나치 함수에 정의된 Docstring이 출력되지 않는다.

help(fibonacci)

# 데코레이션을 사용하지 않은 경우
>>>
Help on function fibonacci in module __main__:

fibonacci(n)
    n번째 피보나치 수를 반환한다.
    
# 데코레이션을 사용한 경우
>>>
Help on function wrapper in module __main__:

wrapper(*args, **kwargs)

 

문제의 원인은 trace 함수가 자신의 본문에 정의된 wrapper 함수를 반환하기 때문이다.

이를 해결하기 위해 functools 내장 모듈에 정의된 wraps 도우미 함수를 사용할 수 있다.

from functools import wraps

def trace(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        print(f'{func.__name__}({args!r}, {kwargs!r}'
              f'-> {result!r}')
        return result
    return wrapper