[Python] decorator와 contextmanager, wraps의 사용

이번 포스팅에서는 파이썬 데코레이터에 대한 개념과 이 데코레이터와 관련된 개념들인 contextmanager, wraps의 쓰임에 대해서도 알아보도록 하자.

 

---

1. decorator란?

데코레이트는 말 그대로 '꾸미다' 라는 의미이다. 하지만 파이썬에서는 꾸미다라는 의미가 직관적이지는 않은 듯 하다. 그래서 개념 확립이 좀 필요할 듯 하다. 우선 데코레이터를 사용하는 예시 코드를 살펴보자.

 

def decorator(func):
    print("decorator init")
    return func

@decorator
def add_ten(x):
    return x + 10

res = add_ten(100)
print(res)

# 출력
# decorator init
# 110

 

decorator 라는 이름의 함수를 만들었다. 우리는 이 함수의 argument로 '함수' 자체를 넣어주고 그 함수를 반환하도록 할 것이다. 위 예시에서는 func 이라는 것을 인자로 받아서 곧바로 func을 리턴하도록 했다. 이제 이것을 데코레이터로 활용할 수 있다.

 

밑에서 @와 활용할 데코레이터로 정의한 decorator 라는 이름을 명시했다. 그리고 add_ten 이라는 함수를 정의하고, 이 함수에는 특정 숫자에 10을 더하도록 하는 로직을 넣었다. 결국, 위에 decorator 라는 이름의 함수 인자로 들어갈 func 이라는 것이 바로 add_ten 이라는 함수가 되는 것이다. 

 

만약 위 로직을 데코레이터를 사용하지 않고 구현하면 어떤 모습이 될까? 그 모습을 보면 데코레이터가 어떤 역할을 대신하는지 파악하기 쉬울 것이다.

 

def decorator(func):
    print("decorator init")
    return func

def add_ten(x):
    return x + 10

deco = decorator(add_ten)
res = deco(100)
print(res)

 

차이점을 보면 decorator 함수를 재할당한 것을 볼 수 있다. 즉, 파이썬의 데코레이터는 이렇게 함수를 재할당하는 과정을 생략하도록 해준다.

 

그런데, 데코레이터를 사용하는 패턴을 보면 데코레이터에 인자를 주는 경우도 있다. 바로 아래처럼 말이다.

 

@decorator("zedd")
def add_ten(x):
    return x + 10

 

이러한 경우는 정의된 decorator 가 함수가 아닌 클래스이고, 그 클래스에서 init 메소드에서 필수적인 argument를 받도록 하고 있기 때문이다. 위와 같이 쓰는 경우의 코드 전문 예시를 보면 아래와 같다.

 

class decorator(object):
    def __init__(self, name):
        print("decorator class init")
        self.name = name

    def __call__(self, func):
        print("decorator class called:", self.name)
        return func

@decorator("zedd")
def add_ten(x):
    return x + 10

res = add_ten(100)
print(res)

# 출력
# decorator class init
# decorator class called: zedd
# 110

 

특이한 점은 위처럼 클래스를 활용해서 데코레이터를 정의해주면 데코레이터 클래스에서 init 메소드와 call 메소드가 순차적으로 실행된다는 것이다. 그리고 난 뒤 데코레이터를 단 실질적인 함수(add_ten 함수)의 로직이 실행된다. 

2. contextmanager

이전에 밑시딥 시리즈를 다루면서 배워보았던 컨텍스트 매니저도 데코레이터 기반으로 동작한다. 그리고 이 컨텍스트 매니저 기반으로 만들어진 함수들 중 우리가 가장 자주 사용하는 함수는 바로 open() 함수이다. 보통 open() 함수를 사용할 때는 with 구문을 사용해서 수행한다. 다시 말해, 컨텍스트매니저 기반으로 만들어진 함수들은 with 구문을 사용해서 동작시킨다. 비슷한 예시로 pytorch의 torch.no_grad() 함수가 존재한다.

 

그러면 컨텍스트 매니저를 우리가 직접 만들어볼 수 있도록 하기 위해서 컨텍스트 매니저를 다루는 방법에 대해 알아보자. 우리만의 컨텍스트 매니저를 만들기 위해서는 아래 처럼 사용해볼 수 있다.

 

from contextlib import contextmanager

@contextmanager
def custom_manager():
    # preprocess

    # main
    yield "Hello world"
    
    # postprocess

 

컨텍스트 매니저를 만들 때는 yield 라는 구문을 기점으로 사전작업, 후처리작업으로 나눌 수 있다. 그리고 yield는 사전작업을 진행하고 메인으로 실행될 작업을 의미한다. 위 코드를 템플릿으로 해서 코드를 좀 더 추가해보자.

 

import os
from contextlib import contextmanager

@contextmanager
def custom_manager():
    # preprocess
    if not os.path.exists("HelloWorld"):
        os.mkdir("HelloWorld")

    # main
    yield "Hello world"
    
    # postprocess
    if os.path.exists("HelloWorld"):
        os.rmdir("HelloWorld")

 

가장 먼저 HelloWorld 라는 디렉토리가 있는지 없는지 확인하고, 없다면 생성하도록 한다. 그리고 메인 작업을 실행하고, 다 끝났으니 후처리 작업으로 이전에 생성했던 HelloWorld 디렉토리를 다시 삭제한다. 이처럼 컨텍스트 매니저는 메인 작업을 실행하기 위해 부가적으로 따라오는 잔잔바리(?) 작업들을 미리 정의해주도록 한다. 이렇게 만들어진 컨텍스트 매니저를 사용하는 프로그래머 입장에서는 단순히 with 구문만을 사용해서 위처럼 HelloWorld 디렉토리가 있는지 체크하고, 생성하고, 삭제하는 작업을 해주지 않아도, 그런 코드를 읽지 않아도 되게 되어 좀 더 고수준의 프로그래밍 경험을 하게 된다.

 

with custom_manager() as f:
    print(f)

 

참고로 위 with 구문에서 작성된 f 라는 값이 바로 yield 로 전달되는 값이 담겨진다.

3. wraps

마지막으로 wraps 에 대해 알아보자. wraps는 functools 라는 파이썬 내장 라이브러리로부터 임포트할 수 있다. wraps의 역할은 간단하게 말해서 데코레이터에 전달된 실질적인 함수를 데코레이터 함수로 포장하지 않고 실질적인 함수로 전달하는 기능을 한다. 당연히 이 텍스트를 보아도 이해가 안갈 것이다. 우선 아래와 같이 wraps를 사용하지 않은 채 작성한 코드부터 살펴보자.

 

def decorator(func):
    """ this is decorator function
    """
    print("decorator init")
    
    def f(*args):
        print("inner f init")
        return func(*args)
        
    return f

@decorator
def add_ten(x):
    """ this is add_ten function
    """
    return x + 10

print(add_ten, add_ten.__doc__)

# 출력
# decorator init
# <function decorator.<locals>.f at 0x1189e9090> None

 

우리는 현재 add_ten 이라는 실질적인 로직을 담당하는 함수를 데코레이터로 감싸주었다. 그런데 출력문에 add_ten 함수의 docstring 내용을 출력했는데, 웬걸 "this is add_ten function" 이라는 문자열이 나올줄 알았는데, None이 나와버렸다. 그리고 함수의 이름도 <function decorator ... > 로 되어있다. 이 말은 add_ten 함수를 출력했지만, 실질적으로는 데코레이터가 껍데기를 감싸고 있다는 것이다.

 

이렇게 데코레이터가 감싸도록 하지 않고 우리의 실질적인 로직을 담당하는 함수를 출력시키고 싶다면 아래처럼 @wraps를 사용할 수 있다.

 

from functools import wraps

def decorator(func):
    """ this is decorator function
    """
    print("decorator init")
    @wraps(func)
    def f(*args):
        print("inner f init")
        return func(*args)
        
    return f

@decorator
def add_ten(x):
    """ this is add_ten function
    """
    return x + 10

print(add_ten, add_ten.__doc__)

# 출력
# decorator init
# <function add_ten at 0x1189e8f70>  this is add_ten function

 

출력을 보면 이제야 함수 이름도 <function add_ten ...> 이라고 나오면 docstring 도 우리가 정의한 "this is add_ten function" 이라는 문자열이 나오는 것을 볼 수 있다.