모던 파이썬 중급 #4 이터러블/제너레이터/yield from

프로그래밍 언어 파이썬 generator iterable
6 분 소요

기초 #4에서 마지막에 잠깐 본 제너레이터 식 (x for x in iter) — 이번 글이 그 주제입니다. for가 어떻게 동작하는지부터 시작해, 사용자 정의 이터러블, 제너레이터 함수, yield from까지 다룹니다.

for in의 정체 — 이터러블 프로토콜 #

우리가 자주 쓰는 
for x in [1, 2, 3]:
    print(x)

이 한 줄이 내부적으로 무엇을 하는지 풀어보면:

실제 흐름 
items = [1, 2, 3]
it = iter(items)        # 1)이터러블 → 이터레이터
while True:
    try:
        x = next(it)    # 2) 다음 값을 요청
    except StopIteration:
        break           # 3) 끝나면 멈춤
    print(x)

핵심 두 단계:

  1. iter(obj) — 이터러블에서 이터레이터를 얻음 (__iter__ 호출)
  2. next(it) — 다음 값을 요청 (__next__ 호출), 끝났으면 StopIteration 던짐

이터러블 vs 이터레이터 #

용어가 헷갈리는 부분이니 정리하겠습니다.

정의메소드예
이터러블 (Iterable)iter()가 가능한 모든 것__iter__list, dict, str, range, 파일, 제너레이터
이터레이터 (Iterator)“다음 값"을 줄 수 있는 것__next__ (와 __iter__)iter([1,2,3])의 결과, 제너레이터

모든 이터레이터는 이터러블 입니다 (자기 자신을 반환하는 __iter__가 있음). 반대는 아닙니다 — list는 이터러블이지만 이터레이터는 아닙니다. next(my_list)는 에러가 납니다.

사용자 정의 이터러블 — 클래스로 #

Range 직접 만들기 
class MyRange:
    def __init__(self, start: int, stop: int):
        self.start = start
        self.stop = stop

    def __iter__(self):
        return MyRangeIterator(self.start, self.stop)

class MyRangeIterator:
    def __init__(self, current: int, stop: int):
        self.current = current
        self.stop = stop

    def __iter__(self):
        return self

    def __next__(self):
        if self.current >= self.stop:
            raise StopIteration
        value = self.current
        self.current += 1
        return value

for x in MyRange(0, 3):
    print(x)
# 0, 1, 2

두 클래스 — 이터러블과 이터레이터를 분리합니다. 이터러블은 여러 번 순회 가능하지만, 이터레이터는 한 번 다 쓰면 끝납니다.

여러 번 순회 가능한 이유 
r = MyRange(0, 3)
list(r)    # [0, 1, 2]
list(r)    # [0, 1, 2]  ← 매번 새 이터레이터

제너레이터 함수 — 같은 일을 한 함수로 #

위 코드 두 클래스를 yield가 들어간 함수 하나로 줄일 수 있습니다.

제너레이터 함수 
def my_range(start: int, stop: int):
    current = start
    while current < stop:
        yield current
        current += 1

for x in my_range(0, 3):
    print(x)
# 0, 1, 2

함수 본문에 yield가 한 번이라도 들어가면, 그 함수는 호출 시 일반 값을 반환하는 게 아니라 제너레이터 객체를 반환합니다. 위 클래스 두 개와 동일한 일을 합니다.

yield가 작동하는 방식 #

가장 헷갈리는 부분이 이것입니다.

제너레이터 만들기 
def gen():
    print("step 1")
    yield 1
    print("step 2")
    yield 2
    print("step 3")

g = gen()
# 여기까지 함수 본문은 아직 실행 안 됨!

g = gen()만으로는 함수 본문이 실행되지 않습니다. 첫 next(g)가 와야 시작합니다.

값 요청 
print(next(g))
# step 1
# 1

print(next(g))
# step 2
# 2

print(next(g))
# step 3
# StopIteration ← yield가 더 없으면

각 yield에서 함수가 일시정지 합니다. 다음 next() 호출이 오면 그 지점부터 재개합니다. 함수가 실행되는 시간이 호출자와 인터리브 되는 게 제너레이터의 핵심입니다.

제너레이터 식과 어떻게 다른가 #

기초 #4에서 본 (x for x in iter)와 같은 동작을 하지만 함수 형태가 더 표현력이 큽니다.

제너레이터 식 vs 함수 
# 한 줄에 표현 가능 — 식이 어울림
squares = (x ** 2 for x in range(10))

# 복잡한 로직 — 함수가 어울림
def squares_evens_only():
    for x in range(10):
        if x % 2 != 0:
            continue
        yield x ** 2

게으름의 가치 — 메모리와 속도 #

제너레이터의 가장 큰 가치는 모든 값을 한 번에 만들지 않는다는 점입니다.

100만 개 처리 
# list 컴프리헨션 — 100만 개 즉시 생성, 메모리 다 사용
squares_list = [x ** 2 for x in range(1_000_000)]

# 제너레이터 — 필요할 때만 생성, 메모리 거의 안 듬
squares_gen = (x ** 2 for x in range(1_000_000))

total = sum(squares_gen)   # 한 번 순회로 끝

무한 시퀀스도 가능 #

무한 시퀀스 
def counter(start: int = 0):
    n = start
    while True:
        yield n
        n += 1

# 처음 5개만 사용
from itertools import islice
first_five = list(islice(counter(), 5))
print(first_five)   # [0, 1, 2, 3, 4]

리스트로는 불가능한 경우입니다. 제너레이터는 요청한 만큼만 값을 만드니까 무한해도 괜찮습니다.

파이프라인 — 제너레이터들을 이어붙이기 #

각 단계가 제너레이터인 데이터 처리 파이프라인을 만들면 메모리 효율이 좋고 의도가 분명합니다.

파이프라인 
def read_lines(path: str):
    with open(path) as f:
        for line in f:
            yield line.rstrip()

def filter_errors(lines):
    for line in lines:
        if "ERROR" in line:
            yield line

def parse_timestamp(lines):
    for line in lines:
        ts, _, msg = line.partition(" ")
        yield (ts, msg)

# 합쳐서 사용
errors = parse_timestamp(filter_errors(read_lines("app.log")))
for ts, msg in errors:
    print(ts, msg)

각 단계가 한 줄씩만 처리 합니다. 100GB 파일이라도 메모리에 다 올리지 않습니다.

yield from — 제너레이터 위임 #

다른 이터러블의 값을 그대로 흘려보내고 싶을 때 쓰는 도구입니다.

🚫 직접 풀어 쓰기 
def chain_two(a, b):
    for x in a:
        yield x
    for y in b:
        yield y
✅ yield from 
def chain_two(a, b):
    yield from a
    yield from b

같은 일을 합니다만 yield from이 더 짧고, 그 외에도 두 가지 추가 효용이 있습니다:

  1. send/throw가 자동 위임됨 (아래에서 다룸)
  2. 하위 제너레이터의 반환값을 받을 수 있음

트리/재귀 순회에 자연 #

재귀로 트리 평탄화 
def flatten(items):
    for item in items:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item

result = list(flatten([1, [2, [3, [4]], 5]]))
print(result)   # [1, 2, 3, 4, 5]

yield from flatten(...) 한 줄이 재귀를 자연스럽게 풀어줍니다.

send, throw, close — 코루틴 기능 #

제너레이터는 값을 받을 수도 있습니다. yield의 결과를 변수에 받으면, 외부에서 send로 값을 넣을 수 있습니다.

send 
def echo():
    while True:
        received = yield
        print(f"받음: {received}")

g = echo()
next(g)            # 첫 yield 까지 진행 (priming)
g.send("hello")    # 받음: hello
g.send("world")    # 받음: world

이 메커니즘 위에 비동기 (#7)와 협력 멀티태스킹이 만들어져 있습니다. 다만 일반 코드에서 send를 직접 다룰 일은 거의 없습니다. 개념만 인지 해 두면 됩니다.

throw와 close 
g.throw(ValueError, "예외 주입")   # 제너레이터 안에서 raise와 동일
g.close()                          # 제너레이터 종료 (GeneratorExit 던짐)

close()는 잘 쓰입니다 — 자원 정리가 필요한 제너레이터에서 try/finally로 정리 코드를 넣으면 close 시점에 실행됩니다.

자원 정리 
def read_lines(path):
    f = open(path)
    try:
        for line in f:
            yield line
    finally:
        f.close()

이 코드는 제너레이터를 끝까지 안 돌려도 (break로 빠져나가도) 가비지 컬렉션 시점에 close()가 호출되어 파일이 닫힙니다.

itertools — 표준 라이브러리의 보석 #

데이터 파이프라인에서 자주 쓰는 도구가 itertools에 있습니다.

자주 쓰는 itertools 
from itertools import (
    count, cycle, repeat,                # 무한
    islice,                               # 슬라이스
    chain,                                # 연결
    groupby,                              # 그룹핑
    accumulate,                           # 누적
    combinations, permutations, product,  # 조합
    starmap, filterfalse, dropwhile, takewhile,  # 변환/필터
)

# 처음 N 개
list(islice(count(), 5))             # [0, 1, 2, 3, 4]

# 여러 이터러블 잇기
list(chain([1, 2], [3, 4]))          # [1, 2, 3, 4]

# 누적합
list(accumulate([1, 2, 3, 4]))       # [1, 3, 6, 10]

# 그룹핑 (정렬 필요)
data = [("a", 1), ("a", 2), ("b", 3)]
for key, group in groupby(data, key=lambda x: x[0]):
    print(key, list(group))
# a [('a', 1), ('a', 2)]
# b [('b', 3)]

데이터 파이프라인을 짤 일이 잦으면 한 번 훑어두면 평생 도움이 됩니다.

표준 라이브러리의 컬렉션도 같은 프로토콜 #

collections.abc의 ABC들은 이 프로토콜을 형식화한 것입니다.

collections.abc — 인터페이스 
from collections.abc import Iterable, Iterator

def consume(items: Iterable[int]) -> int:
    total = 0
    for x in items:
        total += x
    return total

# list, tuple, set, generator, range, ... 모두 통과
consume([1, 2, 3])
consume(range(10))
consume(x for x in [1, 2, 3])

함수가 받는 타입을 **Iterable[T]**로 적으면 가장 넓고 안전합니다. 굳이 list[T]로 좁힐 필요가 없습니다 — 호출자가 제너레이터를 넘기든 셋을 넘기든 똑같이 동작합니다.

@contextmanager가 사실 제너레이터다 #

#3의 @contextmanager가 어떻게 동작하는지 이제 보입니다.

@contextmanager 다시 
from contextlib import contextmanager

@contextmanager
def chdir(path):
    old = os.getcwd()
    os.chdir(path)
    try:
        yield path
    finally:
        os.chdir(old)

이 함수는 yield가 있으니 제너레이터 함수 입니다. @contextmanager가 그 제너레이터를 받아 __enter__가 첫 next를 호출하고, __exit__가 두 번째 next 또는 throw를 호출하는 객체로 감쌉니다. 컨텍스트 매니저가 제너레이터 위에 만들어진 도구라는 점이 여기서 확인됩니다.

정리 #

이번 글에서 잡은 것:

  • for = iter() + next() + StopIteration의 설탕
  • 이터러블 (__iter__) ⊃ 이터레이터 (__iter__ + __next__)
  • 제너레이터 함수 — yield 한 번이라도 있으면 함수 호출이 제너레이터 객체 반환
  • yield마다 일시정지, 다음 next에서 재개
  • 게으름의 가치 — 메모리/무한 시퀀스/파이프라인
  • yield from — 다른 이터러블 위임, 재귀에 자연
  • send/throw/close — 코루틴 메커니즘, close와 try/finally로 자원 정리
  • itertools 표준 도구 모음
  • 함수 인자는 가장 넓게 Iterable[T]로 적자
  • @contextmanager는 제너레이터 위에 얹은 도구

다음 글(#5 데코레이터 패턴)에서는 함수와 클래스를 감싸는 도구 — 데코레이터의 모든 패턴을 다룹니다. @contextmanager, @dataclass도 그 한 형태였습니다.

X