15. 제너레이터 사용하기

💡 제너레이터(generator)란?
이터레이터를 생성해주는 함수로 "발생자"라고 부르기도 한다. 함수 안에서 yield 키워드를 사용함으로써 제너레이터로 동작하는 함수를 만들 수 있다. 제너레이터 객체는 __iter__, __next__ 메소드를 가지고 있으며, 제너레이터 객체에서 __next__ 메소드를 호출할 때마다 함수 안의 마지막 yield까지 코드를 실행하며 yield에서 값을 발생(generate)시킨다. 따라서 이름을 제너레이터라고 한다.

정리해보면 이터레이터를 만드는 방법은 다음 네 가지가 있다.
1. iterable.__iter__()
2. __iter__, __next__ 메소드를 구현한 클래스의 객체 생성
3. __getitem__ 메소드를 구현한 클래스의 객체 생성
4. yield를 사용한 함수(제너레이터)의 객체 생성

 

# 1. 제너레이터(generator)와 yield 알아보기

제너레이터를 만들기 위해서 함수 안에서 "yield 값/변수/함수호출" 형태로 사용한다. 우선 값과 변수를 가지는 제너레이터 예제는 아래와 같다.

# 값을 요소로 가지는 제너레이터 생성
def number_generator():
    yield 0
    yield 1
    yield 2

for num in number_generator():
    print(num)
    
# 실행결과
0
1
2

 

# 변수를 요소로 가지는 제너레이터 생성
a = 10
b = 20
c = 30

def variable_generator():
  yield a
  yield b
  yield c

for var in variable_generator():
  print(var)
  
# 실행결과
10
20
30

 

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

다음과 같이 dir 함수를 사용했을 때 __iter__ 메소드를 확인할 수 있으므로 제너레이터 객체는 이터레이터이다. __next__ 메소드도 있으므로 요소를 차례대로 꺼낼 수도 있다.

g = variable_generator()
print(g)
dir(g)

# 실행 결과
<generator object variable_generator at 0x7f6b783da6d0>
['__class__', '__del__', '__delattr__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', 
 '__getattribute__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__iter__', '__le__', 
 '__lt__', '__name__', '__ne__', '__new__', '__next__', '__qualname__', '__reduce__', 
 '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclasshook__', 'close', 
 'gi_code', 'gi_frame', 'gi_running', 'gi_yieldfrom', 'send', 'throw']

 

실제로 제너레이터 객체의 __next__ 메소드를 세 번 호출하면 각각 10, 20, 30 이 나오고 4번째 호출 시에는 StopIteration 예외가 발생한다.

>>> g.__next__()
10
>>> g.__next__()
20
>>> g.__next__()
30
>>> g.__next__()
---------------------------------------------------------------------------

StopIteration                             Traceback (most recent call last)

<ipython-input-24-079072e7e13f> in <module>
      3 g.__next__()
      4 g.__next__()
----> 5 g.__next__()

StopIteration:

 

이처럼 함수에 yield만 사용해서 간단하게 이터레이터를 구현할 수 있다.  제너레이터 객체와 __iter__, __next__ 메소드로 구현한 클래스 객체간의 차이점은 다음과 같다.

차이점	__iter__, __next__ 메소드로 구현한 클래스 객체	제너레이터 객체
다음 값 반환 방식	__next__ 메소드 안에서 직접 return값으로 반환	yield에 지정한 값을 __next__ 메소드의 return 값으로 반환
StopIteration 예외 발생 방식	__next__ 메소드 안에서 초기값이 끝나는 값보다 커지는 조건에서 직접 raise로 StopIteration 발생시킴	함수의 끝(마지막 yield)까지 도달하면 StopIteration 예외 자동 발생

 

 

1-2. for와 제너레이터

for 반복문을 실행하면 제너레이터 객체의 __iter__ 메소드를 호출해서 제너레이터 객체 자체를 반환하고, 반복할 때마다 제너레이터 객체의 __next__를 호출함으로써 yield에서 발생시킨 값을 가져온다. 그리고 StopIteration 예외가 발생하면 반복을 끝낸다.

 

1-3. yield의 동작 과정 알아보기

 yield는 '생산하다'라는 뜻과 '양보하다'라는 뜻 2개를 가지고 있다. 즉, yield를 실행하면 값을 함수 바깥으로 전달하면서 코드 실행을 함수 바깥에 양보한다. 현재 함수를 잠시 중단하고 함수 바깥의 코드가 실행되도록 만드는 것이다.

 

 예시로 g = number_generator()로 제너레이터 객체 g를 만들고 next(g)를 실행하면 제너레이터 안의 첫 번째 yield가 실행되어 숫자 0을 전달하고 함수 바깥의 코드가 실행되도록 양보한다. 이 때 number_generator 함수는 대기상태가 되고, 함수 바깥에서는 전달받은 값을 a에 할당하여 print(a)를 실행한다. next(g)를 두 번째 실행하면 제너레이터 안의 두 번째 yield가 실행되어 숫자 1을 전달하고 함수 바깥에서는 전달받은 값을 b에 할당하여 print(b)를 실행한다. next(g)를 세 번째 실행하면 세 번째 yield가 실행되어 숫자 2를 전달하고 함수 바깥에서는 전달받은 2를 c에 할당하여 print(c)를 실행함으로써 코드 실행이 끝난다. 만약 next(g)를 한번 더 실행하게 되면 number_generator 함수에서 StopIteration 예외를 발생시킨다.

def number_generator():
    yield 0
    yield 1
    yield 2
    
# 제너레이터 객체(이터레이터) 생성
g = number_generator()

# next 함수를 호출하여 첫번째 yield값을 a에 할당
a = next(g)
print(a)

# next 함수를 두번째로 호출하여 두번째 yield값을 b에 할당
b = next(g)
print(b)

# next 함수를 세번째 호출하여 세번째 yield값을 c에 할당
c = next(g)
print(c)

 

 참고로 return은 반환 즉시 함수가 끝나지만 yield는 잠시 함수 바깥의 코드가 실행되도록 양보하여 값을 가져가게 한 뒤 다시 제너레이터 안의 코드를 계속 실행하는 방식이다. 제너레이터는 함수 끝까지 도달하면 StopIteration 예외가 발생하므로 return을 사용해 함수를 끝내도 StopIteration 예외가 발생한다. 제너레이터 안에서 return에 반환값을 지정하면 StopIteration의 예외 메세지로 들어간다.

def one_generator():
    yield 1
    return "StopIteration 예외 발생"
    
it = one_generator()

try:
    a = next(it)
    print(a)
    b = next(it)
    print(b)
    
except StopIteration as e:
    print(e)
    
# 실행결과
1
StopIteration 예외 발생

 

 

# 2. 값을 여러 번 바깥으로 전달하는 제너레이터 만들기

2-1. while문 안에서 yield  사용하기

def generator(stop):
    n = 0
    while n < stop:
        yield n
        n += 1
        
for i in generator(3):
    print(i)

# 실행결과
0
1
2

 

 

2-2. for문 안에서 yield 사용하기 (feat. yield에서 함수 호출)

yield에 변수를 지정했을 때 변수에 할당된 값이 함수 바깥으로 전달되었던 것처럼 yield 키워드로 함수를 호출해도 결과만 바깥으로 전달한다.

def upper_generator(x):
  for i in x: # x 안의 요소들을 반복하면서 
    yield i.upper() # 해당 요소를 대문자로 만든 결과값을 생산하고 함수 바깥에 전달한다. 

fruits = ['apple', 'banana', 'mango', 'strawberry']

for fruit in upper_generator(fruits): # fruits의 요소 i들을 yield i.upper() -> 변수 fruit에 할당
  print(fruit) # 변수 fruit 출력
  
# 실행결과
APPLE
BANANA
MANGO
STRAWBERRY

 

# 3. yield from으로 값을 여러 번 바깥으로 전달하기

지금까지는 yield로 값을 한 번씩만 함수 바깥으로 전달했다. 그래서 여러 번 값을 전달해야할 때는 while문이나 for문을 사용해 yield를 반복 사용했다. 이런 경우에는 매번 반복문을 사용하지 않고 yield from을 사용하면 된다. yield from에는 반복 가능한 객체, 이터레이터, 제너레이터 객체를 지정한다.

💡 yield from 사용법
- yield from 반복 가능한 객체
- yield from 이터레이터
- yield from 제너레이터

 

3-1. yield from에 반복 가능한 객체 지정하기

yield from을 사용하면  위에서 for문을 활용해 작성한 upper_generator 함수를 for문 없이도 만들 수 있다. yield from은 fruits 리스트 안에 있는 요소를 한개씩 변수 fruit에 전달한다. 즉, 아래 generator 함수에서는 yield from을 한 번 사용하여 값을 네 번 바깥으로 전달하는 것이다. 따라서 next 함수도 네 번 호출할 수 있다.

def generator(x):
    yield from x

fruits = ['apple', 'banana', 'mango', 'strawberry']
for fruit in generator(fruits):
    print(fruit.upper())
  
# 실행결과
APPLE
BANANA
MANGO
STRAWBERRY

 

3-2. yield from에 제너레이터 객체 지정하기

먼저 yield from에 지정할 제너레이터 "number_generator"를 만든다. number_generator는 매개변수로 받은 숫자 직전까지 숫자를 생성하고 함수 바깥으로 해당 숫자를 전달한다. 다음으로는 yield from을 사용한 제너레이터 "rtn_generator"를 만든다. number_generator(3)는 숫자 세 개(0, 1, 2)를 만들어내므로 rtn_generator는 for문 속에서 __next__ 메서드가 호출될 때마다 rtn_generator의 바깥에 0, 1, 2를 한 번에 하나씩 총 세 번 값을 전달한다.

# 매개변수로 받은 숫자 직전까지만 값을 생성하고 해당 값을 전달하는 제너레이터
def number_generator(stop):
  n = 0
  while n < stop:
    yield n
    n += 1

# yield from에 제너레이터 객체를 받는 제너레이터
def rtn_generator(stop): 
  yield from number_generator(stop)

# yield from 한 번으로 지정한 숫자의 횟수만큼 함수 바깥으로 결과값을 전달하는 제너레이터의 객체를 for문의 반복 대상으로 설정
for num in rtn_generator(3):
  print(num)
  
# 실행결과
0
1
2

 

3-3. 제너레이터 표현식

리스트 표현식을 사용할 때에는 []를 사용했다. 같은 리스트 표현식을 ()로 묶으면 제너레이터 표현식이 된다. 리스트 표현식은 처음부터 리스트의 요소를 만들어내지만 제너레이터 표현식은 필요할 때 요소를 만들어내므로 메모리를 절약할 수 있다.

💡 제너레이터 표현식 사용법
(식 for 변수 in 반복 가능한 객체)

# 리스트 표현식
>>> [i for i in range(50) if i % 3 == 1]
[1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49]

# 제너레이터 표현식
>>> (i for i in range(50) if i % 2 == 1)
<generator object <genexpr> at 0x7f6b7838e350>

 

 

예제 1. 파일 읽는 제너레이터 만들기 (본문 40.5 연습문제)

다음 소스 코드에서 words.txt 파일을 한 줄씩 읽은 뒤 내용을 함수 바깥에 전달하는 제너레이터를 작성하세요. 파일의 내용을 출력할 때 파일에서 읽은 \n은 출력되지 않아야 합니다(단어 사이에 줄바꿈이 두 번 일어나면 안 됨).

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

for i in file_read():
  print(i)

# 번외_여러 줄 읽어오기
def file_read():
    with open('words.txt') as file:
        for word in file.readlines():
            yield word.strip('\n')
     
 
for i in file_read():
    print(i)

대용량 데이터를 부분부분 처리할 때 이렇게 제너레이터를 활용한다!

 

 

예제 2. 소수 제너레이터 만들기 (본문 40.6 심사문제)

표준 입력으로 정수 두 개가 입력됩니다(첫 번째 입력 값의 범위는 10~1000, 두 번째 입력 값의 범위는 100~1000이며 첫 번째 입력 값은 두 번째 입력 값보다 항상 작습니다). 다음 소스 코드에서 첫 번째 정수부터 두 번째 정수 사이의 소수(prime number)를 생성하는 제너레이터를 만드세요. 소수는 1과 자기자신만으로 나누어 떨어지는 1보다 큰 양의 정수입니다.

# 
def is_prime_number(num):
  for n in range(2, ((num // 2) + 1)):
    if num % n == 0:
      return False
  return True

def prime_number_generator(start, stop):
  for i in range(start, stop):
    if is_prime_number(i) == True:
      yield i

start, stop = map(int, input().split())
 
g = prime_number_generator(start, stop)
print(type(g))
for i in g:
    print(i, end=' ')

# "에라토스네테스의 체" 방식을 이용하여 소수를 찾기
def prime_number_generator(start, stop):
    for num in range(start, stop):
        is_prime_num = True
        for i in range(2, num): # num이 포함되면 아무 숫자도 남지 않으므로 주의
            if num % i == 0:
                is_prime_num = False
        if is_prime_num == True:
            yield num
            
start, stop = map(int, input().split())

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