[가짜연구소 3기] 데이터 엔지니어링 47 - More on Decorators

주제

좀 더 유용한 데코레이터를 작성하는 방법에 대해 배웠다.

---

예제 살펴보기

import time

def timer(func):
	"""함수 수행 시간 출력하는 데코레이터
	Args:
		func (callable): 데코레이팅 된 함수
	
	Returns:
		callable: 데코레이팅 된 함수
	"""
	def wrapper(*args, **kwargs):
		t_start = time.time()
		result = func(*args, **kwargs)
		t_total = time.time() - t_start
		print('{} took {}s'.format(func.__name__, t_total))
		return result

	return wrapper

@timer
def sleep_n_seconds(n):
	time.sleep(n)

sleep_n_seconds(5)

sleep_n_seconds took 5.01214075088501s

---

def memoize(func):
	"""
	데코레이팅된 함수의 결과를 빠르게 찾기 위해 캐시에 저장
	"""
	cache = {}
	def wrapper(*args, **kwargs):
# kwargs가 dict이라서 계속 에러남 -> mutable type은 dict key 불가
		if (args, kwargs) not in cache:
			cache[(args, kwargs)] = func(*args, **kwargs)
		return cache[(args, kwargs)]
	return wrapper

@memoize
def slow_function(a, b):
	print('Sleeping...')
	time.sleep(10)
	return a + b

slow_function(3, 4)

import time

def memoize(func):
	"""
	데코레이팅된 함수의 결과를 빠르게 찾기 위해 캐시에 저장
	"""
	cache = {}
	def wrapper(*args, **kwargs):
		key = str(args)+str(kwargs)
		if key not in cache:
			cache[key] = func(*args, **kwargs)
		return cache[key]
	return wrapper

@memoize
def slow_function(a, b):
	print('Sleeping...')
	time.sleep(10)
	return a + b

slow_function(3, 4)

연산결과를 cache에 저장했기 때문에 바로 출력이 된다.

그리고 여러 함수에 타이머 함수를 사용하고 싶다면 반복하지 말고 데코레이터를 사용하는 것이 좋다.

---

데코레이터의 단점

데코레이터의 문제점 : 데코레이팅 된 함수의 메타 데이터를 가리게 된다.

__doc__으로 독스트링을 확인하거나 함수 이름을 찾기 위해 __name__ 등을 사용했을 때, 정보를 제대로 얻을 수 없게 된다.

import time

def timer(func):
	"""함수 수행 시간 출력하는 데코레이터
	Args:
		func (callable): 데코레이팅 된 함수
	
	Returns:
		callable: 데코레이팅 된 함수
	"""
	def wrapper(*args, **kwargs):
		t_start = time.time()
		result = func(*args, **kwargs)
		t_total = time.time() - t_start
		print('{} took {}s'.format(func.__name__, t_total))
		return result

	return wrapper

여기서 중첩된 함수는 wrapper()이다.

__doc__을 사용하면 wrapper() 함수에 작성한 doc이 출력된다.

def add_hello(func):
  # Add a docstring to wrapper
  def wrapper(*args, **kwargs):
    """Print 'hello' and then call the decorated function."""
    print('Hello')
    return func(*args, **kwargs)
  return wrapper

@add_hello
def print_sum(a, b):
  """Adds two numbers and prints the sum"""
  print(a + b)
  
print_sum(10, 20)
print_sum_docstring = print_sum.__doc__
print(print_sum_docstring)

Hello
30
Print 'hello' and then call the decorated function.

functools의 wraps를 사용해서 이 문제를 해결할 수 있다.

from functools import wraps
def timer(func):
	"""함수 수행 시간 출력하는 데코레이터
	Args:
		func (callable): 데코레이팅 된 함수
	
	Returns:
		callable: 데코레이팅 된 함수
	"""
	@wraps(func)
	def wrapper(*args, **kwargs):
		t_start = time.time()
		result = func(*args, **kwargs)
		t_total = time.time() - t_start
		print('{} took {}s'.format(func.__name__, t_total))
		return result

	return wrapper

@timer
def sleep_n_seconds(n):
	time.sleep(n)

sleep_n_seconds(5)

sleep_n_seconds.__name__

sleep_n_seconds took 5.005099058151245s
sleep_n_seconds

원래의 함수를 부를 때는 __wrapped__ 로 호출할 수 있다.

---

데코레이터에 인수 넣기 - 데코레이터 팩토리

def run_three_times(func):
	def wrapper(*args, **kwargs):
		for i in range(3):
			func(*args, **kwargs)
	return wrapper

위 데코레이터는 반드시 3번 실행된다.

반복 횟수를 사용자의 입력을 인수로 전달받아 설정하고 싶지만 데코레이터는 func 만 가질 수 있다.

그리고 데코레이터는 괄호를 가질 수 없다.

기존 데코레이터를 감싸는 데코레이터 함수를 만들어서 해결할 수 있다.

def run_n_times(n):
	def decorator(func):
		def wrapper(*args, **kwargs):
			for i in range(n):
				func(*args, **kwargs)
		return wrapper
	return decorator

@run_n_times(3)
def print_sum(a, b):
	print(a + b)

run_five_times = run_n_times(5)

@run_five_times
def print_sum(a, b):
  print(a + b)
  
print_sum(4, 100)

# Modify the print() function to always run 20 times
print = run_n_times(20)(print)

print('What is happening?!?!')

---

Timeout()

함수가 예상보다 오래걸리면 에러를 발생시키는 데코레이터 예제

from functools import wraps
import signal

def raise_timeout(*args, **kwargs):
    raise TimeoutError()

signal.signal(signalnum=signal.SIGALRM, handler=raise_timeout)
signal.alarm(5) #5초 알람 설정
signal.alarm(0) #알람 취소

def timeout_in_5s(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        signal.alarm(5)
        try :
            return func(*args, **kwargs)
        finally:
            signal.alarm(0)
    return wrapper


import time

@timeout_in_5s
def foo():
    time.sleep(10)
    print('foo!')

foo()

TimeoutError

유연한 timeout 데코레이터 만들기

def timeout(n):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            signal.alarm(n)
            try:
                return func(*args, **kwargs)
            finally:
                signal.alarm(0)
        return wrapper
    return decorator


@timeout(5)
def foo():
    time.sleep(10)
    print('foo!')

@timeout(20)
def bar():
    time.sleep(10)
    print('bar!')


foo() # TimeoutError
bar() # bar!

Uploaded by Notion2Tistory v1.1.0