[python cleancode] 7. 제너레이터 사용하기

출처

파이썬 클린코드를 읽으며 정리한 내용입니다.

• 제너레이터는 전통적인 언어와 파이썬을 구분짓는 또 다른 특징적인 기능

제너레이터 만들기

• 한번에 하나씩 구성요소를 반환해주는 이터러블을 생성하는 객체
• 고성능이면서도 메모리를 적게 사용하는 반복을 위한 방법
  • 특정 요소를 어떻게 만드는지 아는 객체를 만들어서 필요할 때마다 하나씩만 가져온다!
  • 게으른 연산으로 무거운 객체를 사용하도록 함

제너레이터 개요

• 대규모의 구매 정보에서 최저 판매가, 최고 판매가, 평균 판매가를 구하는 예시
  • 모든 구매정보를 파일에서 읽어들여 반환한다면 데이터의 양에 따라 시간이 오래 걸리고, 메모리가 부족할수도 있음
  • 한번에 하나의 데이터만 사용한다면 제너레이터를 사용하는 것이 좋다

def load_purchases(filename):
    with open(filename) as f:
        for line in f:
            *_, price_raw = line.partition(",")
            yield float(price_raw)

• 결과를 담을 리스트가 없어지고, return 문 또한 사라진다.
• 파이썬에서 어떤 함수라도 yield를 사용하면 제너레이터 함수가 된다.
• 이렇게 리스트 대신 제너레이터를 사용해도, 구매정보를 하나씩 가져와 지표를 계산하는 코드는 수정이 필요없다.
• 제너레이터 객체는 이터러블이며, 이터러블을 사용하면 for 루프의 다형성을 보장하는 강력한 추상화가 가능하다.

제너레이터 표현식

• 제너레이터도 리스트, 셋, 사전처럼 컴프리헨션으로 정의될 수 있다.

# list comprehension
[x**2 for x in range(10)]

>>> [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

# generator expression(comprehension)
(x**2 for x in range(10))

>>> <generator object <genexpr> at 0x7f92b8447e40>

• min(), max, sum()같은 이터러블 연산에는 제너레이터 표현식을 전달하는 것이 파이썬스럽다.

이상적인 반복

• 파이썬에서 반복할 때 유용하게 사용할 수 있는 관용적인 코드
• 이러한 코드를 보며 제너레이터를 활용하는 방법을 익혀보자.

관용적인 반복 코드

• 내장 함수인 enumerate와 유사한 객체를 만들어보자.

# 시작 값을 입력하면 무한 시퀀스를 만드는 클래스
class NumberSequence:
    def __init__(self, start=0):
        self.current = start
        
    def next(self):
        current = self.current
        self.current += 1
        return current

• 위 인터페이스에 기반을 두어 클라이언트를 작성하면 명시적으로 next()함수를 호출해야 한다.

seq = NumberSequence()
seq.next()

>>> 0

seq.next()

>>> 1

• 이 객체를 반복문에서 사용하려면 __iter__매직 메서드를 구현해 객체가 반복 가능하게 만들어야 한다.
• 여기에 더해, next()메서드를 수정하여 __next__ 매직 매서드를 구현하면 객체는 이터레이터가 된다.

class SequenceOfNumbers:
    def __init__(self, start=0):
        self.current = start
    
    def __next__(self):
        current = self.current
        self.current += 1
        return current
    
    def __iter__(self):
        return self

• 이렇게 하면 요소를 반복할수도 있고, __next__ 메서드를 구현했으므로 next()내장함수를 사용할 수 있다.

list(zip(SequenceOfNumbers(), "abcdef"))

>>> [(0, 'a'), (1, 'b'), (2, 'c'), (3, 'd'), (4, 'e'), (5, 'f')]

seq = SequenceOfNumbers(100)
next(seq)

>>> 100

next(seq)

>>> 101

next()함수

• next()내장 함수는 이터러블을 다음 요소로 이동시키고 기존의 값을 반환한다.

word = iter("hello")
next(word)

>>> 'h'

next(word)

>>> 'e'

• 이터레이터가 더이상 값을 가지고 있지 않다면 StopIteration 예외가 발생한다.

next(word)
next(word)
next(word)
next(word)

>>> 
StopIterationTraceback (most recent call last)

<ipython-input-15-a62e3e37f15b> in <module>
      2 next(word)
      3 next(word)
----> 4 next(word)

StopIteration: 

• 이 문제를 해결하고 싶다면 예외를 캐치하거나, next()함수의 두번째 파라미터에 기본값을 제공할 수 있다.

next(word, "default")

>>> 'default'

제너레이터 사용하기

• 앞의 코드는 제너레이터를 사용해 클래스를 사용하지 않고도 훨씬 간단하게 작성할 수 있다.

def sequence(start=0):
    while True:
        yield start
        start += 1

• 제너레이터는 위와 같이 무한루프를 사용해도 완벽하게 안전하다.
• 제너레이터 함수가 호출되면 yield 전까지 실행되고, 값을 생성한 후 그 자리에서 멈춘다.

seq = sequence(10)
next(seq)

>>> 10

next(seq)

>>> 11

itertools

• 이터레이션이 언어의 중요한 컴포넌트이기 때문에, 이터러블로 작업하면 코드가 파이썬 자체와 더 잘 어울리게 된다.
• itertools 모듈을 사용하면 그 기능을 온전히 활용할 수 있다.
• 이터레이터, 제너레이터, itertools를 연결하면 새로운 객체를 만들수도 있다.
• 처음 예제로 돌아가서, 지표 계산 시 특정 기준을 넘은 값에 대해서만 연산을 하려면?
  • while문 안에 조건을 추가할수도 있지만, 파이썬스럽지 못하며 수정사항을 잘 반영할 수 없다.
  • itertools를 이용하면 코드를 유지하며 요구사항을 반영할 수 있다.

from itertools import islice
purchases = islice(filter(lambda p: p > 1000.0, purchases), 10)

이터레이터를 사용한 코드 간소화

• 여러번 반복하기
  • 이제 이터레이터에 대해 배웠으므로 처음의 예제를 더욱 간소화할 수 있다.

def process_purchases(purchases):
    # 원래의 이터러블을 세 개의 새로운 이터러블로 분신술
    min_, max_, avg = itertools.tee(purchases, 3)
    return min(min_), max(max_), median(avg)

• 중첩 루프
  • 경우에 따라 1차원 이상을 반복해서 값을 찾아야 하는 경우 제너레이터가 유용할 수 있다.
  • 두 단계를 벗어나야 하는 경우 플래그나 예외를 사용하지 말자..

def _iterate_array2d(array2d):
    for i, row in enumerate(array2d):
        for j, cell in enumerate(row):
            yield (i, j), cell
            
def search_nested(array, desired_value):
    try:
        coord = next(
            coord for (coord, cell) in _iterate_array2d(array)
            if cell == desired_value
        )
    except StopIteration:
        raise ValueError(f"{desired_value} not found")
        
    print("%i, %i에서 값 %r 찾음" % (*coord, desired_value))
    return coord

• 위와 같이 보조 제너레이터를 이용해 반복을 추상화하면 더 많은 차원의 배열을 사용하는 경우에도 클라이언트는 그것에 대해 알 필요가 없이 기존 코드를 사용하면 된다.
• 이것이 이터레이터 디자인 패턴의 본질이다.

파이썬의 이터레이터 패턴

• 파이썬에서의 반복을 더 자세히 이해하기 위해 제너레이터로부터 약간 벗어나보자.
• 훌륭한 이터러블 객체를 만들게 되면 효율적이고 컴팩트한 코드를 작성할 수 있다.
• 이터러블은 __iter__와 __next__ 매직 메서드를 구현한 객체이다.
  • 일반적으로 이렇게 구현하지만, 항상 이 두가지를 모두 구현해야 하는 것은 아니다.
  • __iter__ 구현 -> 이터러블
  • __next__ 구현 -> 이터레이터

이터레이션 인터페이스

• 이터러블은 반복할 수 있는 어떤 것으로, 실제로 반복할 때에는 이터레이터를 사용한다.
• 즉 __iter__를 통해 이터레이터를 반환하고, __next__를 통해 반복 로직을 구현한다.

# 이터러블하지 않은 이터레이터 객체의 예시
class SequenceIterator:
    def __init__(self, start=0, step=1):
        self.current = start
        self.step = step
    
    def __next__(self):
        value = self.current
        self.current += self.step
        return value

si = SequenceIterator(1, 2)
next(si)

>>> 1

next(si)

>>> 3

# 시퀀스에서 하나씩 값을 가져올 수는 있지만 반복할 수 없다.
for _ in SequenceIterator(): pass

>>> 
TypeErrorTraceback (most recent call last)

<ipython-input-32-3e62347f344e> in <module>
      1 # 시퀀스에서 하나씩 값을 가져올 수는 있지만 반복할 수 없다.
----> 2 for _ in SequenceIterator(): pass

TypeError: 'SequenceIterator' object is not iterable

이터러블이 가능한 시퀀스 객체

• __iter__를 구현한 객체는 for 루프에서 사용할 수 있지만, 반복을 위해 이것을 구현해야 하는 것은 아니다.
• 객체가 시퀀스인 경우, 즉 __getitem__과 __len__을 구현한 경우에도 반복 가능하다.
  • 이 경우 인터프리터는 IndexError 예외가 발생할 때까지 순서대로 값을 제공한다.

class MappedRange:
    """특정 숫자 범위에 대해 맵으로 변환"""
    def __init__(self, transformation, start, end):
        self._transformation = transformation
        self._wrapped = range(start, end)
        
    def __getitem__(self, index):
        value = self._wrapped.__getitem__(index)
        result = self._transformation(value)
        print(f"index {index}: {result}")
        return result
    
    def __len__(self):
        return len(self._wrapped)

mr = MappedRange(abs, -10, 5)
mr[0]

>>> 
index 0: 10

10

mr[-1]

index -1: 4

4

list(mr)

>>> 
index 0: 10
index 1: 9
index 2: 8
index 3: 7
index 4: 6
index 5: 5
index 6: 4
index 7: 3
index 8: 2
index 9: 1
index 10: 0
index 11: 1
index 12: 2
index 13: 3
index 14: 4

[10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 1, 2, 3, 4]

• 객체가 시퀀스여서 반복이 가능할수는 있지만, 기본적으로 반복을 위한 객체를 디자인할 때는 __iter__ 매직 메서드를 구현해 정식 이터러블 객체를 만드는 것이 좋다.

코루틴

• __iter__와 __next__를 구현한 제너레이터 객체를 이용해 단순히 다음 요소로 반복, 이동할수도 있지만 이런 기본 기능 외에도 제너레이터를 코루틴으로 활용할수도 있다.
• 코루틴을 지원하기 위해 추가된 기본 메서드는 아래와 같다.
  • .close()
  • .throw(ex_type[, ex_value[, ex_traceback]])
  • .send(value)

제너레이터 인터페이스의 메서드

• 앞에서 언급한 각 메서드에 대해 알아본 후, 코루틴을 만들어보자.

close()

• 이 메서드를 호출하면 제너레이터에서 GeneratorExit 예외가 발생한다.
• 종료 상태를 지정하는데에 사용되는 예외
• 코루틴이 자원 관리를 하는 경우 이 예외를 통해 코루틴이 보유한 모든 자원을 해제할 수 있다.
• 컨텍스트 관리자나 finally 블록과 비슷하지만 더 명확하다.

# 코루틴 예시
def stream_db_records(db_handler):
    try:
        while True:
            yield db_handler.read_n_records(10)
    except GeneratorExit:
        db_handler.close()

• 위 예시에서는 제너레이터를 호출할 때마다 데이터베이스 핸들러에서 얻은 레코드를 10개씩 반환하고, close()를 호출하면 데이터베이스 연결도 함께 종료한다.

throw(ex_type[, ex_value[, ex_traceback]])

• 현재 제너레이터가 중단된 위치에서 예외를 던진다. 제너레이터가 예외를 처리했으면 제너레이터 내 except절에 있는 코드가 호출되고, 예외를 처리하지 않았으면 예외가 호출자에게 전파된다.

class CustomException(Exception):
    pass

def stream_data(db_handler):
    while True:
        try:
            yield db_handler.read_n_records(10)
        except CustomException as e:
            logger.warning("처리 가능한 에러 %r, 계속 진행", e)
        except Exception as e:
            logger.error("처리할 수 없는 에러 %r, 중단", e)
            db_handler.close()
            break

• streamer.throw(CustomException)이 발생한 경우 제너레이터는 INFO 레벨의 메시지를 기록하고 다음 yield 구문으로 이동한다.
• except Exception as e 구문이 없다면 streamer.throw(RuntimeError)이 발생한 경우 제너레이터가 중지된 행에서 예외가 호출자에게 전파되고, 제너레이터는 중지된다.

send(value)

• 지금까지 만든 제너레이터가 읽어올 레코드의 개수를 파라미터로 받아서 호출하도록 수정해보자.

def stream_db_records(db_handler):
    retrieved_data = None
    previous_page_size = 10
    try:
        while True:
            page_size = yield retrieved_data
            if page_size is None:
                page_size = previous_page_size
            
            previous_page_size = page_size
            retrieved_data = db_handler.read_n_records(page_size)
    except GeneratorExit:
        db_handler.close()

• 위와 같이 코루틴을 작성하면 send() 메서드를 통해 인자 값을 전달할 수 있다. send() 메서드를 사용하면 할당 구문 오른쪽에 yield 키워드가 나오게 되고, 할당 구문의 변수를 사용해 인자 값을 사용할 수 있다.
• 이 경우 yield 키워드는 두 가지 일을 한다.
  • retrieved_data를 호출자에게 보내고 그 곳에서 멈춘다.
  • send() 메서드를 통해 전달된 produced 값을 받는다.
• 코루틴에 값을 전송하는 것은 yield 구문이 멈춘 상태에서만 가능하기 때문에, 코루틴을 yield까지 이동시켜야 한다. 이를 위해 next()를 적어도 한 번은 호출해야 한다.
• 이 메서드는 제너레이터와 코루틴을 구분하는 기준이 된다.

# 더 간단한 버전
def stream_db_records(db_handler):
    retrieved_data = None
    previous_page_size = 10
    try:
        while True:
            page_size = (yield retrieved_data) or page_size
            retrieved_data = db_handler.read_n_records(page_size)
    except GeneratorExit:
        db_handler.close()

• next()를 반드시 호출해야 한다는 것을 기억할 필요 없이 코루틴을 생성하자마자 바로 사용할 수 있다면 훨씬 편할 것이다. 이를 위해 데코레이터를 사용해 코루틴을 초기화할 수 있다.

# 데코레이터 사용 예시
@prepare_coroutine
def stream_db_records(db_handler):
    retrieved_data = None
    previous_page_size = 10
    try:
        while True:
            page_size = (yield retrieved_data) or page_size
            retrieved_data = db_handler.read_n_records(page_size)
    except GeneratorExit:
        db_handler.close()

코루틴 고급 주제

• 코루틴은 충분히 멋지지만, 실제로 복잡한 코루틴을 병렬로 실행하려면 추가 기능이 필요하다.
• 각종 문제를 해결하기 위해 yield from이라는 생성자 구문이 도입되었다.

코루틴에서 값 반환하기

• 제너레이터에서는 단순히 한 번에 하나씩 값을 생성하기만 하면 되기 때문에, for 루프의 모든 단계에서 생성되는 각 값에 대해서만 신경쓰면 된다.
• 반면 코루틴은 일반적으로 반복보다는 상태를 중단하는 데에 초점을 맞추기 때문에, 경량 스레드라고 생각할 수 있고 이 경우 값을 반환하는 것이 이해된다.
• 하지만 제너레이터에서 값을 반환하면 반복이 중단된다.

def generator():
    yield 1
    yield 2
    return 3

value = generator()
next(value)

>>> 1

next(value)

>>> 2

try:
    next(value)
except StopIteration as e:
    print(f"returned value: {e.value}")

>>> returned value: 3

작은 코루틴에 위임하기 - yield from 구문

• 코루틴(제너레이터)이 값을 반환할 수 있다는 점은 흥미롭지만, 언어에서 기능을 지원해주지 않으면 귀찮은 면이 있다.
• 이것을 개선하기 위해 yield from이 등장했다.
• yield from을 사용하는 가장 간단한 예시로는 제너레이터 체인이 있다.
  • 제너레이터 체인은 여러 제너레이터를 하나의 제너레이터로 합치는 기능을 하며, 서브 제너레이터의 값을 한 번에 수집할 수 있게 해준다.

# 여러 이터러블을 받아 하나의 스트림으로 반환하는 에시
def chain(*iterables):
    for it in iterables:
        for value in it:
            yield value

• 여기에서 yield from을 사용하면 서브 제너레이터에서 직접 값을 생산할 수 있으므로 중첩 루프를 피할 수 있다.

def chain(*iterables):
    for it in iterables:
        yield from it

list(chain("hello", ["world"], ("tuple", "of", "values")))

>>> ['h', 'e', 'l', 'l', 'o', 'world', 'tuple', 'of', 'values']

• yield from 구문은 어떤 이터러블에 대해서도 동작하며 이것을 사용하면 최상위 제너레이터가 직접 값을 yield한 것과 같은 효과를 나타낸다.
• 하지만 이것이 yield from의 존재 이유는 아니다.

서브 제너레이터에서 반환된 값 구하기

def sequence(name, start, end):
    print("%s 제너레이터 %i에서 시작" % (name, start))
    yield from range(start, end)
    print("%s 제너레이터 %i에서 종료" % (name, end))
    return end

def main():
    step1 = yield from sequence("first", 0, 5)
    step2 = yield from sequence("second", step1, 10)
    return step1 + step2

g = main()
next(g)

>>> 
first 제너레이터 0에서 시작

0

next(g)

1

next(g)

>>> 
first 제너레이터 5에서 종료
second 제너레이터 5에서 시작

5

next(g)

>>> 
second 제너레이터 10에서 종료

StopIterationTraceback (most recent call last)

<ipython-input-77-e734f8aca5ac> in <module>
----> 1 next(g)

StopIteration: 15

• 위 예제와 같이 yield from을 사용하면 코루틴의 종료 시 최종 반환 값을 구할 수 있다.
• 헷갈리는 문법
  • a = yield b: send()를 통해 받은 값을 a에 할당하고 b를 호출자에게 보낸 후 멈춘다.
  • a = yield from b: b 이터러블에서 값을 하나씩 받아와 yield하고 b가 최종적으로 반환한 값을 a에 할당한다.
  • b 이터러블(서브 제너레이터)에서 반환하는 값이 없는 경우 yield from의 반환값은 None이다.

서브 제너레이터와 데이터 송수신하기

• 코루틴의 진정한 강력함을 느낄 수 있게 해주는 멋진 기능
• 코루틴 역할을 하는 제너레이터는, 값을 전송하고 예외를 던지면 값을 받아서 처리하거나 반드시 예외를 처리해야 하며, 서브 제너레이터에 위임한 코루틴에 대해서도 마찬가지이다.
• 수동으로 이러한 것들을 처리하면 매우 복잡하다.

def sequence(name, start, end):
    value = start
    print("%s 제너레이터 %i에서 시작" % (name, value))
    while value < end:
        try:
            received = yield value
            print("%s 제너레이터 %r값 수신" % (name, received))
            value += 1
        except CustomException as e:
            print("%s 제너레이터 %s 에러 처리" % (name, e))
            received = yield "OK"
    return end

def main():
    step1 = yield from sequence("first", 0, 5)
    step2 = yield from sequence("second", step1, 10)
    return step1 + step2

g = main()
next(g)

>>> 
first 제너레이터 0에서 시작

0

g.send("first generator")

>>> 
first 제너레이터 'first generator'값 수신

1

g.throw(CustomException("custom exception"))

>>> 
first 제너레이터 custom exception 에러 처리

'OK'

• sequence 서브 제너레이터에 값을 보내지 않고도 yield from을 통해 sequence에 데이터를 전달할 수 있다.
• 첫 번째 코루틴이 멈춰진 상태에서 데이터를 전송하거나 예외를 던지면 첫 번째 코루틴 인스턴스가 값을 받는다.
• 첫 번째 코루틴이 끝나면 step1 변수와 두 번째 코루틴에 end값이 전달된다.
• 각 코루틴이 생성하는 값 또한 main 제너레이터에 전파된다.

비동기 프로그래밍

• 지금까지 살펴본 것들을 활용해 비동기 프로그램을 만들 수 있다.
• 이러한 기능을 통해 얻을 수 있는 가장 큰 장점은 논블로킹 방식으로 병렬 I/O 작업을 할 수 있다는 것이다.
  • 보통 서드파티 라이브러리에서 구현한 저수준의 제너레이터를 이용해 코루틴이 일시 중단된 동안 실제 I/O 처리를 한다.
  • 코루틴이 정지된 동안 프로그램은 다른 작업을 할 수 있어 효율적이다.
• 코루틴과 제너레이터가 기술적으로 동일하다는 점에서 혼란스럽지만, 의미적으로 다르다.
  • 효율적인 방법을 원할 때는 제너레이터를 사용하고, 논블로킹 I/O 작업을 원할 때는 코루틴을 사용한다.
• 코루틴과 제너레이터를 혼합해서 사용하다가 런타임 오류가 발생하기도 한다. 혼동하지 않도록 코루틴을 명시적으로 타이핑하는것이 좋다.
• 새로운 구문으로 await와 async def를 이용할수도 있다.
  • await는 yield from을 대신하기 위한 용도로 awaitable 객체에 대해서만 동작하며, 코루틴은 awaitable이다.
  • async def는 @coroutine 데코레이터를 대신하여 코루틴을 정의하는 방법이다. 이것은 호출 시 실제로 객체를 만들어 코루틴 인스턴스를 반환한다.
  • 한마디로 yield from은 코루틴을 기다리는 old way, await는 modern way
    • 참고
  • asyncio는 기존 제너레이터 기반 코루틴도 지원한다.
  • 업데이트된 내용을 급하게 포함한 느낌..

# 새로운 코루틴 작성 방법
import asyncio
import time

async def say_after(delay, what):
    await asyncio.sleep(delay)
    print(what)

async def main():
    print(f"started at {time.strftime('%X')}")

    await say_after(1, 'hello')
    await say_after(2, 'world')

    print(f"finished at {time.strftime('%X')}")

# for plain python
# asyncio.run(main())

# for jupyter - 이미 주피터 노트북이 이벤트 루프를 수행중이므로 만들어주지 않아도 됨
await main()

started at 13:14:31
hello
world
finished at 13:14:34

• 파이썬에서 비동기 프로그래밍을 한다는 것은 일련의 코루틴을 관리하는 이벤트 루프(asyncio 등)가 있다는 뜻이다.
  • 일련의 코루틴들은 이벤트 루프에 속하며, 이벤트 루프의 스케줄링 메커니즘에 따라 호출된다.
  • 각 코루틴이 실행되면 사용자가 작성한 내부 코드가 실행되고, 다시 이벤트 루프에 제어권을 반납하려면 await coroutine을 호출한다.
• 참고: 코루틴과 태스크