[python] Decorator(데코레이터) - 02

안녕하세요. TDR입니다.

오늘은 저번에 이어 class를 이용한 decorator 구현법과 decorator에 변수를 할당하는 방법을 알아보겠습니다.

(Decorator - 01은 아래 링크 참조)

[python] Decorator(데코레이터) - 01

이번에는 평소와는 다르게 실용적인 예시를 통해 class를 이용한 decorator & Decorator에 변수 전달하기를 같이 봐보겠습니다.

import time

time_dict = dict()

class StartTime:
    """
    함수를 호출한 시간을 기록해 놓는 decorator class
    """
    def __init__(self, name):
        self.name = name
  
    def __call__(self, func):
        def wrapper(*args):
            time_dict[self.name] = time.time()
            func(*args)
        return wrapper
		

class EndTime:
    """
    StartTime decorator class에서 time_dict에 넣어 놓은 시간과 비교하여 출력
    """
    def __init__(self, name):
        self.name = name
		
    def __call__(self, func):
        def wrapper(*args):
            func(*args)
            start = time_dict.pop(self.name, -1)
            if start == -1:
                print(f'{self.name} is not exist!')
            else:
                end = time.time()
                print(f'{self.name} : {end - start:.3f} sec')
        return wrapper


class CheckTime:
    """
    함수의 총 실행 시간을 출력해주는 decorator class
    """
    def __init__(self, func):
        self.func = func

    def __call__(self):
        start = time.time()
        self.func()
        end = time.time()
        print(f'{self.func.__name__} : {end - start:.3f} sec')


@StartTime('apple')
def func_01():
    time.sleep(1)

@EndTime('banana')
def func_02():
    time.sleep(2)

@StartTime('banana')
def func_03():
    time.sleep(3)

@EndTime('apple')
def func_04():
    time.sleep(4)

@CheckTime
def all_func_call():
    func_01()
    func_02()
    func_03()
    func_04()

func_01()
func_03()
func_04()
func_02()
## Result
# apple : 8.008 sec
# banana : 9.009 sec

all_func_call()
## Result
# banana is not exist!
# apple : 10.011 sec
# all_func_call : 10.011 sec

코드를 간략히 설명하면, StartTime, EndTime, CheckTime이라는 3개의 class와 해당 클래스를 decorator로 쓰는 func_01~04까지의 함수들이 있습니다. StartTime decorator는 호출되면 시간을 time_dict라는 dictionary에 저장해 놓았다가 EndTime decorator가 호출되면 StartTime이 불린 이후부터의 총 시간을 출력합니다. (dictionary가 잘 생각 안나시는 분은 아래 링크 참조)

[python] Dictionary

하지만 잘 보면 class decorator로 'apple', 'banana'와 같은 인수를 전달하고 있습니다. 이로인해 단순히 하나의 함수의 시간을 측정하는 것이 아닌, 원하는 구간의 시간을 측정할 수 있습니다.

__init__ 함수를 보면, 인수로 self와 name이 있습니다. self는 class의 함수에 필수로 들어가야 하는 값이고 name은 decorator에서 매개변수로 넣은 값임을 알 수 있습니다.

__call__함수는 class를 함수와 같이 호출 가능하게 해주는 클래스 고유 함수입니다. 이를통해 class를 마치 function과 같이 호출 할 수 있게 해줍니다. 인수로 self와 func가 있습니다. func는 decorator를 붙인 함수를 의미합니다. 내부 함수인 wrapper 함수에서 func(*args)를 통해 decorator가 붙은 함수를 호출해 준 다음 time_dict에서 name key값이 있다면 StartTime에서 저장해 놓은 시간을 가져옵니다. (없으면 예외처리) 그리고는 총 소요 시간을 출력하게 됩니다.

마지막으로 CheckTime class를 보면 위 클래스들과는 조금 다른데, __init__의 인수로 func를 받아 self에 저장하고, __call__에서는 별도의 인수를 받지 않고, self.func를 호출하기 앞뒤로 시간을 측정해서 func가 동작하는데 걸린 시간만을 계산해서 출력해주고 있습니다. 이 decorator는 호출한 단일 함수의 시간을 측정하기 위한 것입니다.

해당 코드는 실제로 제가 코드 작성 시 사용하는 코드를 좀 더 간략하게 수정한 것으로, 위와 같이 decorator를 통해 기존 코드 수정 없이 특정 기능을 추가할 수 있습니다.