[Python] 제너레이터 generator

제너레이터란?

제너레이터는 이터레이터를 생성해주는 함수이다.

이터레이너틑 클래스에 __iter__, __next__ 또는 __getitem__ 메서드를 구현해야 하지만 제너레이터는 함수 안에서 yield 키워드만 사용하면 끝이다.

제너레이터와 yield

함수 안에서 yield를 사용하면 제너레이터가 되며 yield에는 값(변수)를 지정한다.

제너레이터 만들기

def number_generator():
    yield 0
    yield 1
    yield 2


for i in number_generator():
    print(i)

0
1
2

제너레이터 객체가 이터레이터인지 확인하기

def number_generator():
    yield 0
    yield 1
    yield 2


## 제너레이터 객체가 이터레이터인지 확인
g = number_generator()
print(g)
print(dir(g))

함수에 yield만 사용해서 간단하게 이터레이터를 구현할 수 있다.

이터레이터는 __next__ 메서드 안에서 직접 return으로 값을 반환했지만 제너레이터는 yield에 지정한 값이 __next__ 메서드의 반환값으로 나온다.

또한 이터레이터는 raise로 StopIteration 예외를 직접 발생시켰지만 제너레이터는 함수의 끝까지 도달하면 StopIteration 예외가 자동으로 발생한다.

제너레이터는 제너레이터 객체에서 __next__ 메서드를 호출할 떄마다 함수 안의 yield까지 코드를 실행하며 yield에서 값을 발생(generaate) 시킨다. 그래서 이름이 제너레이터(generator)인 것이다.

yield의 동작 과정 알아보기

def number_generator():
    yield 0  # 0을 함수 바깥으로 전달하면서 코드 실행을 바깥에 양보
    yield 1  # 1을 함수 바깥으로 전달하면서 코드 실행을 바깥에 양보
    yield 2  # 2을 함수 바깥으로 전달하면서 코드 실행을 바깥에 양보


g = number_generator()

a = next(g)  # yield를 사용하여 함수 바깥으로 전달한 값은 next의 반환값으로 나온다
print(a)

b = next(g)
print(b)

c = next(g)
print(c)

0
1
2

yield를 사용하여 바깥으로 전달한 값은 next 함수(__next__ 메서드)의 반환값으로 나온다. 즉, 제너레이터 함수가 실행되는 중간에 next로 값을 가져온다.

이렇게 제너레이터는 함수를 끝내지 않은 상태에서 yield를 사용하여 값을 바깥으로 전달할 수 있다.

return은 반환 즉시 함수가 끝나지만 yield는 잠시 함수 바깥의 코드가 실행되도록 양보하여 값을 자겨가게 한 뒤 다시 제너레이터 안의 코드를 계속 실행하는 방식이다.

제너레이터 만들기

def number_generator(stop):
    n = 0
    while n < stop:
        yield n
        n += 1


for i in number_generator(3):
    print(i, end=" "

0
1
2

yield에서 함수 호출하기

다음과 같이 yield에서 함수를 호출할 수도 있다.

def upper_generator(x):
    for i in x:
        yield i.upper()  # 함수의 반환 값을 바깥으로 전달


fruits = ["apple", "pear", "grape", "pineapple", "orange"]
for i in upper_generator(fruits):
    print(i)

APPLE
PEAR
GRAPE
PINEAPPLE
ORANGE

yield는 무엇을 지정하든 결과만 바깥으로 전달한다.

yield from

반복문을 사용하지 않고 반복 가능한 객체에서 값을 하나씩 바깥으로 전달할 수 있게 해주는 것이 yield from이다.

def number_generator():
    x = [1, 2, 3]  # 반복 가능한 객체
    yield from x


for i in number_generator():
    print(i)

1
2
3

yield from에 제너레이터 객체 지정하기

yield from에 제너레이터 객체를 지정해줄 수도 있다.

def number_generator(stop): # 제너레이터
    n = 0
    while n < stop:
        yield n
        n += 1


def three_generator():
    yield from number_generator(3) # 제너레이터 객체 지정


for i in three_generator():
    print(i)

0
1
2

제너레이터 표현식

리스트 표현식을 괄호 ( )로 묶으면 제너레이터 표현식이 된다.

리스트 표현식은 처음부터 리스트의 요소를 만들어내지만 제너레이터 표현식은 필요할 때 요소를 만들어내므로 메모리를 절약할 수 있다.

x = (i for i in range(50) if i % 2 == 0)

print(x)

for i in x:
    print(i, end=" ")

예제 1 - 파일 읽기 제너레이터 만들기

def file_read():
    with open("words.txt") as file:
        while True:
            line = file.readline()
            if line == "":
                break
            yield line.strip("\n")


for i in file_read():
    print(i)

예제 2 - 소수 제너레이터 만들기

def prime_number_generator(start, stop):
    n = start
    while n < stop:
        flag = 0
        for i in range(2, n // 2):
            if n % i == 0:
                flag = 1
                break
        if flag == 0:
            yield n
        n += 1

## 혹은 다음과 같이 만들 수 있다.
def prime_number_generator2(start, stop):
    for n in range(start, stop):
        is_prime = True
        for i in range(2, n):
            if n % i == 0:
                is_prime = False
        if is_prime == True:
            yield n

start, stop = map(int, input().split())

g = prime_number_generator(start, stop)
print(type(g))
for i in g:
    print(i, end=" ")

참고 자료

• 파이썬 코딩 도장
• SCHOOL OF WEB