파이썬의 generator와 iterator (간단하게)

이름이 비슷해서 그런지 매번 헷갈리는 generator와 iterator!

Iterator

Iterator가 필요한 이유

메모리에 리스트를 만들어 iterate 하는 방법은 메모리를 많이 소모하게 된다. 예를 들면

sum = 0
for i in [1, 2, ..., 10000000]:
  sum += i

과 같이 계산하려면 리스트에 모든 숫자를 미리 할당해둬야 하기 때문이다. 리스트 자체는 sum 을 계산하기 위한 중간 단계기 때문에 메모리에 저장할 필요가 전혀 없다. 그래서 파이썬을 처음 배울 때 반복문을 다음과 같이 구현하도록 배운다.

sum = 0
for i in range(10000000):
  sum += i

여기서 range가 바로 Iterator로, 메모리에 숫자를 미리 저장하지 않고 iteration이 될 때마다 다음 숫자를 꺼내어주는 역할을 한다.

Iterator 이해하기

Iterator는 다음과 같이 두 가지 메소드를 가진 객체로 정의된다.

class Iterator:
  def __iter__():
    return Iterator
  def __next__():
    return next_value

즉 현재 단계에서는 __iter__() 메소드로 다음 단계의 Iterator를 받고 __next__() 메소드로 값을 받아오는 것이다. 1씩 증가하는 간단한 예제를 확인해보자.

class Iterator(object):
    def __init__(self):
        self.index = 0

    def __next__(self):
        if self.index == 10:
            raise StopIteration
        self.index += 1
        return self.index


class Counter(object):
    def __iter__(self):
        iter = Iterator()
        return iter
    
a = Counter()
print([i for i in a])

# [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

객체에 __iter__ 메소드가 있으면 Iterator로 작동하는 원리이다. 다만 StopIteration 이 정의되어 있지 않으면 동작하지 않는다.

Generator

Generator는 Iterator와 느낌은 비슷한데 동작은 완전히 다르다. 이는 yield 키워드 기반의 lightweight coroutine이다. 아래 예제를 실행시켜 보자.

def test1():
    print("print 1")
    yield 1
    print("print 2")
    yield 2


def test2():
    # Simpler iterator implementation
    for i in range(10):
        yield i

test1() # Do nothing
t1 = test1()
print("Type of t1 :", type(t1))
print(dir(t1))
next(t1)
next(t1)
# next(t1) <-- This will raise StopIteration exeption

출력

Type of t1 : <class 'generator'>
['__class__', '__del__', '__delattr__', ... 'send', 'throw']
print 1
print 2

함수 test1을 실행시켜도 아무 응답이 없다. 이 함수를 next 로 호출해야 작동하게 되는데, 동작이 iterator와 동일하다. t1 객체를 보면 __next__가 있어서 실제로 iterator처럼 동작함을 알 수 있다. 이때 next로 인해 실행되는 부분은 yield 키워드를 만날 때까지로, 매번 next가 호출될 때마다 다음 yield키워드까지 실행되며, 만일 yield가 없다면 StopIteration 예외가 발생한다.

Generator를 더 쉽게 만드는 방법은 다음과 같다.

(i for i in range(10))
# <generator object <genexpr> at 0x7fcd604b9190>

Reference

https://blog.humminglab.io/python-coroutine-programming-1/