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

파이썬의 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/