[Python] 제너레이터

제너레이터를 사용하는 이유

제너레이터는 왜 사용하는 걸까요? 제너레이터를 사용하면 이터러블(iterable) 객체를 생성하여 큰 데이터 집합의 구성요소를 한 번에 하나씩 반환하게 하여, 메모리를 절약할 수 있습니다. 따라서 좀 더 유연하고 메모리 효율적인 코드를 작성할 수 있게 되죠.

제너레이터는 이터레이터이므로 먼저 이터레이터에 대해 간단히 알아본 후, 제너레이터의 사용 예시에 대해 알아보겠습니다.

 

---

이터레이터(iterator)

이터레이터는 next 함수 호출 시 계속해서 그 다음 값을 리턴해주는 객체입니다.

 

리스트를 iter 함수로 감싼 후 간단한 이터레이터를 만들 수 있습니다.

lst = [1, 2, 3]
lst = iter(lst)

next(lst)	# 1
next(lst)	# 2
next(lst)	# 3
next(lst)	# StopIteration

next() 내장 함수는 이터레이터를 다음 요소로 이동시키고 기존의 값을 반환하는 함수인데요, 이처럼 next 함수 호출 마다 이터레이터의 요소값들이 차례대로 리턴되고 한계치를 넘으면 StopIteration 에러가 발생하게 됩니다. 이터레이터는 next로 그 값을 읽어 들이면 다시 그 값을 읽을 수 없는 특징이 있습니다.

 

이와 같은 이터레이터의 특성은 클래스 내부의 __next__라는 메소드에 의해 나타는건데요, 클래스로 이터레이터를 만드는 객체를 구현해보겠습니다.

class MyIterator:
    def __init__(self, number):
        self.number = number
        self.count = 0
        
    def __next__(self):
        self.count += 1
        return self.count + self.number
        
obj = MyIterator(10)
next(obj)	# 11
next(obj)	# 12
next(obj)	# 13

 

그러나 위처럼 만든 이터레이터는 이터러블(iterable)하지 않습니다. for 루프로 이터레이터를 작동시킬 수가 없다는 얘기죠. 

이터레이터를 이터러블하게 만들어 주기 위해선 내부에 __iter__메소드를 만들어 줘야 합니다.

class MyIterator:
    def __init__(self, number):
        self.number = number
        self.count = 0
        
    def __iter__(self):
        return self
        
    def __next__(self):
        self.count += 1
        return self.count + self.number
        
obj = MyIterator(10)
for i in obj:
    print(i)	# 11 12 13 14 ...

 

정리하자면, 객체를 이터러블하게 만들어주는 메소드는 __iter__고, 이터레이터로 만들어주는 메소드는 __next__인 것입니다. 그리고 앞서 말한 이터레이터의 특성대로, for문을 이용하여 요소값을 반복하고 난 후에는 다시 반복하더라도 더 이상 그 값을 다시 가져오지 못합니다.

 

---

제너레이터(generator)

이제 제너레이터에 대해 알아봅시다. 제너레이터는 함수가 마치 이터레이터처럼 하나의 값만을 리턴하는게 아닌 연속된 값들을 순차적으로 리턴할 수 있게 해줍니다. 

제너레이터를 사용하려면 클래스를 만들지 않고, 단지 return 대신 yield를 사용하는 함수를 만들면 됩니다.

def gen(start):
    while True:
        yield start
        start += 1
    
g = gen(1)
next(g)		# 1
next(g)		# 2
next(g)		# 3

 

제너레이터 특성을 사용한 함수 내에서 yield 키워드를 만나게되면, 그 값을 리턴하고 현재 상태를 기억한 채로 다른 값을 요청할 때까지(다른 yield문에 도달할 때까지) 실행을 일시 중지합니다. 따라서 위 코드처럼 무한 루프를 사용해도 안전한 것이죠.

이처럼 모든 값을 한번에 계산하지 않고 필요할 때마다 하나의 항목씩 수행하는 제너레이터의 특성을 "느긋한 계산법(Lazy evaluation)" 이라고 부릅니다.

 

그렇다면 이제 제너레이터를 통해 큰 데이터 파일에서 작업에 필요한 데이터만 가져와 메모리를 절약할 수 있는 코드를 작성해보겠습니다.

def read_in_chunks(file_handler, chunksize=1024):
    while True:
        data = file_handler.read(chunksize)
        if not data:
            break
        yield data
        
f = open("large_data.dat")
for piece in read_in_chunks(f):
    print(piece)

 

함수를 사용하지 않고 좀 더 간단하게 제너레이터를 만들 수도 있습니다. 제너레이터 표현식이라 하여 작성된 이 구문은 리스트 컴프리핸션과 비슷하지만, 리스트 대신 튜플을 이용한다는 점이 다릅니다.

gen = (i * i for i in range(1, 1000))

gen = mygen()
print(next(gen))	# 1
print(next(gen))	# 4
print(next(gen))	# 9

 

yield from는 다른 제너레이터에서 값을 얻는 역할을 합니다. 이 키워드를 이용하면 다음과 같이 멀티 리스트를 펼치는 코드를 작성할 수도 있습니다.

def flat_list(lter_values):
    for item in iter_values:
        if hasattr(item, "__iter__"):
            yield from flat_list(item)
        else:
            yield item
            
list(flat_list([1, [2], [3, [4]]]))	# [1, 2, 3, 4]

 

---

참고 자료

• 파이썬 라이브러리 - 이터레이터와 제너레이터 : https://wikidocs.net/134909
• <클린 파이썬 : 효과적인 파이썬 코딩 기법, 수닐 카필 지음>
• <파이썬 클린 코드 : 유지보수가 쉬운 파이썬 코드를 만드는 비결, 마리아노 아나야 지음>