[python] 반복자(iterator)와 제너레이터(generator)

반복자(iterator)를 사용하는 이유

파이썬을 사용하다보면 어느 순간 심심치 않게 iter() 또는 next() 등을 목격했으리라 생각한다. 

 

아마 파이썬으로 처음 프로그래밍을 공부했다면 (필자와 같은 비전공자) "결국에는 for문을 활용해서 사용하는 거 같은데.. 쉽게 그냥 list에 담아두고 사용하면 되지 왜 번거롭게 iterator를 만들 당최 왜 사용할까?" 라는 생각이 들 수도 있다.

 

iterator를 사용하는데 여러 이유가 있겠지만, 기본적으로 시간과 비용 측면에서 매우 효율적이기 때문이다. 

가장 직관적인 예제는 빅데이터를 활용한 딥러닝 모델을 구현할 때인데, 예를 들면 데이터 100만개를 학습한다고 가정해보자.

100만개를 시작부터 불러놓고 학습을 하는 것과 iterator가 필요한 데이터만 즉석으로 주어 학습을 하는 것과 어떤 것이 효율적일까?

당연히 후자다! 

데이터를 한번에 다 부르지 않고, 조금씩 나누어 불러서 사용하는 방식이 유리하다. (이걸 custom dataset이라고 부름)

 

반복자를 활용한 for 문 코드 구현

아래 코드를 보자. 리스트를 for문에 넣어 순회하여 요소를 순서대로 꺼낸다.

for num in [11, 22, 33, 44, 55]:
    print(num)

 

이 반복자는 꽤 쉬운 문장이지만, 코드가 동작하는 내부 과정은 나름 복잡하다.

아래는 실제 for문을 직접 구현한 것이다.

nums = [11, 22, 33, 44, 55]

it = iter(nums)
while True:
    try:
        num = next(it)
    except StopIteration:
        break
    print(num)
    
    
    
# 실행결과
# 11
# 22
# 33
# 44
# 55

iter함수로 nums 리스트의 반복자를 구한다. 이후 while 문으로 무한 루프를 돌며 next 함수로 반복자의 다음 요소를 순서대로 꺼낸다.

컨테이너의 끝에 도달하면 StopIteration 예외가 발생하며 이때 루프를 탈출하게 된다.

 

for 문의 내부 동작을 정확히 이해하고 이 조건만 맞춰 준다면, 아래 코드와 같이 임의의 클래스를 반복 가능 객체로 만들 수 있다.

 

# 이 코드는 점프 투 파이썬에서 발췌했습니다.

class MyItertor:
    def __init__(self, data):
        self.data = data
        self.position = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.position >= len(self.data):
            raise StopIteration
        result = self.data[self.position]
        self.position += 1
        return result
        
        
if __name__ == "__main__":
    i = MyItertor([1,2,3])
    for item in i:
        print(item)

# 출력물
# 1
# 2
# 3

"__iter__ 메서드는 이터레이터 객체를 반환하는 메서드이므로 MyIterator 클래스에 의해 생성되는 객체를 의미하는 self를 반환하도록 했다. __next__ 메서드는 next() 함수 호출 시 수행되므로 MyIterator 객체 생성 시 전달한 데이터를 하나씩 반환하도록 하고 더는 반환할 값이 없으면 StopIteration 예외를 발생시키도록 구현했다." (점프투 파이썬)

제너레이터 (generator)

제너레이터는 이터레이터를 생성하주는 함수라고 생각하면 되겠다.

반복자의 내부 구조는 명확하게 공개되어 있고, 구조도 간단해서 이해하기 쉬운 편이나, 매번 관련 메서드를 작성하는 것은 무척 번거로울 수 있다.

이때 제너레이터를 활용하면 반복자를 더 간단하게 만들 수 있따.

 

제너레이터의 가장 큰 특징은 일반적인 함수 형태를 띠지만 yield 명령으로 값 리턴한다는 점이다.

 

def my_generator():
    for i in range(1, 1000):
        result = i
        yield result

g = my_generator()

print(next(g))  # 1
print(next(g))  # 2
print(next(g))  # 3
print(next(g))  # 4
# ...

이렇게 제너레이터를 활용하면서, 복잡한 이터레이터 클래스 대신 간단한 이터레이터의 경우 쉽게 구현할 수 있다.

 

제너레이터 표현식도 존재하는데, 위 코드보다 더 간단하게 구현할 수 있다. 핵심은 튜플로 for문을 감싸주기만 하면 된다

g = (i for i in range(1, 1000))

 

범위가 좁을 때는 별 차이가 없겠지만, 아주 큰 범위를 다룰 때는 모두 만들어 두고 시작하는 것보다 그때그때 만드는 것이 더 빠르고 메모리 효율도 좋다. 특히 중간에 반복을 그만둘 때는 두 코드의 속도 차이가 현격하게 벌어진다고 한다.