[Python] 데코레이터 만들기

파이썬에서 데코레이터란 @staticmethod와 같은 형태로 클래스나 함수 위에 선언하여 사용하는 것을 말합니다. 데코레이터는 기존 함수 내부의 코드를 수정하지 않고 시작과 끝에 추가 기능을 구현할 수 있습니다.

class Test:
    @staticmethod
    def test():
        pass

또한 데코레이터는 1개 이상 추가할 수도 있습니다.

from abc import abstractmethod

class Test:
    @staticmethod
    @abstractmethod
    def test():
        pass

아래에서는 함수와 클래스로 데코레이터를 만드는 방법을 소개합니다.

함수로 데코레이터 만들기

함수로 데코레이터를 만드는 방법은 아래와 같습니다.

def deco1(func):
    def wrapper(*args, **kwargs):
        print(func.__name__, 'start')
        print(f'args: {args}, kwargs: {kwargs}')
        func(*args, **kwargs)
        print(func.__name__, 'end')

    return wrapper

@deco1
def test(n, val):
    print('test 함수')
    print(n, val)

test(10, val=False)

# test start
# args: (10,), kwargs: {'val': False}
# test 함수
# 10 False
# test end

print(test)
# <function deco1.<locals>.wrapper at 0x7f9418035af0>

간단하게는 위와 같이 만들 수 있지만 test 함수를 출력하면 test 함수명이 나오는 것이 아닌 데코레이터의 함수명이 나오게 됩니다. deco1 함수는 그 안에 정의된 wrapper를 반환합니다. 이 wrapper 함수가 데코레이터를 호출한 후 해당 호출을 담고 있는 모듈의 test라는 이름에 할당되는 값입니다. 이러한 동작은 객체 내부를 조사할 때 문제가 될 수 있으므로 functools의 wraps함수를 사용해야 합니다.

from functools import wraps

def deco1(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print(func.__name__, 'start')
        print(f'args: {args}, kwargs: {kwargs}')
        func(*args, **kwargs)
        print(func.__name__, 'end')

    return wrapper

@deco1
def test(n, val):
    print('test 함수')
    print(n, val)

test(10, val=False)
# test start
# args: (10,), kwargs: {'val': False}
# test 함수
# 10 False
# test end

print(test)
# <function test at 0x7fa7381cfaf0>

데코레이터의 기능은 그대로 보장하면서 의도한 함수명을 출력하도록 도와줍니다. wraps는 데코레이터를 작성하는 데 이용하는 함수(데코레이터)입니다. 이 함수(데코레이터)를 wrapper 함수에 적용하면 내부 함수에 있는 중요한 메타 데이터가 모두 외부함수로 복사됩니다.

여기서 직접 만든 데코레이터에 변수를 추가하려면 함수를 한 번 더 감싸줘야 합니다.

def deco1(name):
    def deco1_decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            print(func.__name__, 'start')
            print(f'{name}: name')
            print(f'args: {args}, kwargs: {kwargs}')
            func(*args, **kwargs)
            print(func.__name__, 'end')

        return wrapper
    return deco1_decorator

@deco1(name='이름')
def test(n, val):
    print('test 함수')
    print(n, val)

test(10, val=False)

# test start
# 이름: name
# args: (10,), kwargs: {'val': False}
# test 함수
# 10 False
# test end

위와 같이 함수를 감싸 3중 함수로 만들면 데코레이터에 파라미터를 전달할 수 있습니다. 위와 같은 형태일 때, @deco1 과 같이 괄호가 빠지게 되면 오류가 발생합니다.

이제 위의 호출 순서를 알아보도록 합니다.

from functools import wraps

print(1)
def deco1(name='이름'):
    print(3, name)
    def deco1_decorator(func):
        print(5, func)
        @wraps(func)
        def wrapper(*args, **kwargs):
            print(8, args, kwargs)
            func(*args, **kwargs)
            print(10)
        print(6)
        return wrapper
    print(4)
    return deco1_decorator

print(2)
@deco1()
def test(n, val):
    print(9, 'test')

print(7)
test(10, val=False)
print(11)

# 1
# 2
# 3 이름
# 4
# 5 <function test at 0x7fbd000c5af0>
# 6
# 7
# 8 (10,) {'val': False}
# 9 test
# 10
# 11

위의 흐름은 아래와 같이 데코레이터를 사용하지 않았을 때와 동일합니다.

print(1)
def deco1(name='이름'):
    print(3, name)
    def deco1_decorator(func):
        print(5, func)
        def wrapper(*args, **kwargs):
            print(8, args, kwargs)
            func(*args, **kwargs)
            print(10)
        print(6)
        return wrapper
    print(4)
    return deco1_decorator

def test(n, val):
    print(9, 'test')

print(2)
deco1_decorator = deco1('이름22')
print('#')
deco1_call = deco1_decorator(test)
print(7)
deco1_call(10, val=False)
print(11)

# 1
# 2
# 3 이름22
# 4
# #
# 5 <function test at 0x7fb23800faf0>
# 6
# 7
# 8 (10,) {'val': False}
# 9 test
# 10
# 11

이처럼 만든 데코레이터는 전처리로 유효성 검사나 입력된 파라미터의 값을 변경하고 후처리로 로그를 쌓을 수 있습니다.

from functools import wraps

def deco1(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        name = (args and args[0]) or kwargs.get('name', None)
        if isinstance(name, str):
            name = name.upper()
        else:
            raise ValueError
        func(name)
        log()
    return wrapper

@deco1
def test(name):
    print(name)

test()

# KIM

클래스로 데코레이터 만들기

클래스는 함수로 만드는 것 보다 더 간단합니다.

from functools import wraps

class Deco1:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        print(self.func.__name__, 'start')
        print(f'args: {args}, kwargs: {kwargs}')
        self.func(*args, **kwargs)
        print(self.func.__name__, 'end')

@Deco1
def test(name):
    print(name)

test('이름')
# test start
# args: ('이름',), kwargs: {}
# 이름
# test end

print(test)
# <__main__.Deco1 object at 0x7fc2900a4280>

클래스로 만들면 단점 중에 하나가 함수로 만들었을 때 @wraps를 붙였던 것처럼 위의 코드에선 붙일 수가 없습니다. 데코레이터에 파라미터가 있으면 쉽게 해결이 가능하지만 위처럼 파라미터가 없는 경우, 아래처럼 구현하고 사용해야 합니다.

from functools import update_wrapper

class Deco1:
    def __init__(self, func):
        self.func = func
        update_wrapper(self, func)

    def __call__(self, *args, **kwargs):
        print(self.func.__name__, 'start')
        print(f'args: {args}, kwargs: {kwargs}')
        self.func(*args, **kwargs)
        print(self.func.__name__, 'end')

@Deco1
def test(name):
    print(name)

test('이름')
print(test)
# <__main__.Deco1 object at 0x7fc2900a4280>
print(test.__name__)
# test

위 형태를 봤을 땐 파라미터를 받지 않는 데코레이터를 구현할 때엔 함수를 사용하는 것이 더 좋아보입니다.

다음은 파라미터를 받는 데코레이터를 구현한 예제입니다.

from functools import wraps

class Deco1:
    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs

    def __call__(self, func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            print(func.__name__, 'start')
            print(self.args, self.kwargs)
            print(f'args: {args}, kwargs: {kwargs}')
            func(*args, **kwargs)
            print(func.__name__, 'end')

        return wrapper

@Deco1("파라미터", asid=123)
def test(name):
    print(name)

test('이름')
# test start
# ('파라미터',) {'asid': 123}
# args: ('이름',), kwargs: {}
# 이름
# test end

print(test)
# <function test at 0x7f8e100f2040>

요약

파라미터가 없는 데코레이터를 만들 땐, 클래스보단 함수로 진행하는 것이 좋다.

데코레이터 생성 후, @wraps는 필수로 붙여주자