제너레이터 사용하기

1. 제너레이터란?

제너레이터는 이터레이터를 생성해주는 함수입니다. 함수 안에서 yield 키워드를 사용하면 그 함수는 제너레이터가 됩니다.

# yield를 사용하면 제너레이터 함수가 됨
def number_generator():
    yield 1
    yield 2
    yield 3

gen = number_generator()
print(gen)  # <generator object number_generator at 0x...>

일반 함수처럼 호출해도 바로 실행되지 않고 제너레이터 객체가 반환됩니다. next()를 호출할 때 비로소 코드가 실행됩니다.

2. yield의 동작 원리

yield는 "양보하다"라는 뜻입니다. 값을 함수 바깥으로 전달하면서 코드 실행을 잠시 중단하고, 함수 바깥의 코드가 실행되도록 양보합니다. 이후 next()가 다시 호출되면 멈췄던 자리부터 이어서 실행됩니다.

return과 비교하면 차이가 더 명확합니다.

def number_generator():
    print("첫 번째 yield 전")
    yield 1                     # 여기서 중단, 1 반환
    print("두 번째 yield 전")
    yield 2                     # 여기서 중단, 2 반환
    print("세 번째 yield 전")
    yield 3                     # 여기서 중단, 3 반환

gen = number_generator()

print(next(gen))  # 첫 번째 yield 전 → 1
print(next(gen))  # 두 번째 yield 전 → 2
print(next(gen))  # 세 번째 yield 전 → 3
print(next(gen))  # StopIteration 예외 발생

3. 이터레이터와 비교

이전 글에서 __iter__와 __next__를 직접 구현했던 Counter 클래스를 제너레이터로 바꾸면 얼마나 간결해지는지 비교해보겠습니다.

# 이터레이터 방식 — 클래스로 구현
class Counter:
    def __init__(self, stop):
        self.current = 1
        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

# 제너레이터 방식 — 훨씬 간결
def counter(stop):
    for i in range(1, stop + 1):
        yield i

for num in counter(3):
    print(num)
# 1
# 2
# 3

같은 동작을 훨씬 적은 코드로 표현할 수 있습니다.

4. yield에서 함수 호출하기

yield 뒤에 함수를 호출하면 해당 함수의 반환값이 바깥으로 전달됩니다.

def get_name():
    return input("이름을 입력하세요: ")

def name_generator():
    yield get_name()   # get_name()의 반환값을 바깥으로 전달
    yield get_name()
    yield get_name()

gen = name_generator()
print(next(gen))  # 입력값 출력
print(next(gen))  # 입력값 출력

5. yield from

yield from을 사용하면 반복 가능한 객체, 이터레이터, 제너레이터 객체에서 값을 하나씩 꺼내 바깥으로 전달할 수 있습니다. 중첩 반복문 없이 간결하게 표현할 수 있습니다.

# yield로 하나씩 전달하는 방식
def generator_a():
    for i in [1, 2, 3]:
        yield i

# yield from으로 간결하게
def generator_b():
    yield from [1, 2, 3]   # 반복 가능한 객체

print(list(generator_a()))  # [1, 2, 3]
print(list(generator_b()))  # [1, 2, 3]

제너레이터 안에서 다른 제너레이터를 연결할 때도 활용할 수 있습니다.

def generator_1():
    yield from [1, 2, 3]

def generator_2():
    yield from [4, 5, 6]

def combined():
    yield from generator_1()   # 제너레이터 연결
    yield from generator_2()

print(list(combined()))  # [1, 2, 3, 4, 5, 6]