이터레이터 사용하기

1. 이터레이터란?

이터레이터(Iterator) 는 값을 차례대로 꺼낼 수 있는 객체입니다. 파이썬에서는 이터레이터를 만들어두고 값이 필요한 시점이 되었을 때 비로소 값을 생성합니다. 이처럼 데이터 생성을 뒤로 미루는 방식을 지연 평가(Lazy Evaluation) 라고 합니다.

덕분에 데이터가 아무리 많아도 한꺼번에 메모리에 올리지 않아도 됩니다. range(1000000)이 메모리를 거의 차지하지 않는 것도 이 때문입니다.

2. 반복 가능한 객체 vs 이터레이터

헷갈리기 쉬운 개념이라 먼저 구분하고 넘어가겠습니다.

dir() 함수로 객체가 어떤 메서드를 가졌는지 확인할 수 있습니다.

# 반복 가능한 객체인지 확인
print('__iter__' in dir([1, 2, 3]))   # True
print('__iter__' in dir(range(10)))   # True

# 이터레이터인지 확인 (__next__ 까지 있어야 함)
it = iter([1, 2, 3])
print('__next__' in dir(it))          # True
print('__next__' in dir([1, 2, 3]))   # False ← 리스트는 이터레이터가 아님

리스트, 문자열, 딕셔너리, 세트, range는 모두 반복 가능한 객체이지만 이터레이터는 아닙니다. iter()를 호출하면 이터레이터를 얻을 수 있습니다.

3. 이터레이터 동작 원리

이터레이터는 __next__로 요소를 하나씩 꺼내다가 더 이상 꺼낼 요소가 없으면 StopIteration 예외를 발생시켜 반복을 끝냅니다.

# iter()로 이터레이터 생성
numbers = [1, 2, 3]
it = iter(numbers)

# next()로 값을 하나씩 꺼냄
print(next(it))  # 1
print(next(it))  # 2
print(next(it))  # 3
print(next(it))  # StopIteration 예외 발생

for 문이 내부적으로 하는 일이 바로 이것입니다. iter()로 이터레이터를 만들고, StopIteration이 발생할 때까지 next()를 반복 호출합니다.

__iter__와 __next__를 모두 지원하는 것을 이터레이터 프로토콜(Iterator Protocol)을 지원한다고 합니다.

4. iter()와 next()의 추가 기능

iter() — 센티넬 값 지정

iter()에 호출 가능한 객체와 반복을 끝낼 값(센티넬, Sentinel) 을 함께 전달하면, 해당 값이 나올 때 자동으로 반복을 끝냅니다.

import random

# random.randint(1, 6)을 호출하다가 6이 나오면 반복 종료
it = iter(lambda: random.randint(1, 6), 6)

for num in it:
    print(num)  # 6이 나오기 전까지 출력

next() — 기본값 지정

next()에 기본값을 지정하면 반복이 끝났을 때 StopIteration 대신 기본값을 반환합니다.

it = iter([1, 2])

print(next(it, 0))  # 1
print(next(it, 0))  # 2
print(next(it, 0))  # 0 ← 꺼낼 값이 없으므로 기본값 반환

5. 이터레이터 언패킹

이터레이터는 언패킹이 가능합니다. 자주 사용하는 map()도 이터레이터이기 때문에 변수 여러 개에 값을 바로 할당할 수 있습니다.

# map도 이터레이터 — 언패킹 가능
a, b, c = map(int, ["1", "2", "3"])
print(a, b, c)  # 1 2 3

6. 이터레이터 직접 만들기

이터레이터를 직접 만들 때는 __iter__와 __next__ 메서드를 구현하거나, __getitem__ 메서드를 구현하면 됩니다.

__iter__ + __next__ 방식

# 1부터 stop까지 값을 차례로 반환하는 이터레이터
class Counter:
    def __init__(self, stop):
        self.current = 1
        self.stop = stop

    def __iter__(self):
        return self         # 이터레이터 자기 자신을 반환

    def __next__(self):
        if self.current > self.stop:
            raise StopIteration  # 반복 종료
        value = self.current
        self.current += 1
        return value

for num in Counter(3):
    print(num)
# 1
# 2
# 3

__getitem__ 방식

# 인덱스로 접근하는 방식 — 파이썬이 자동으로 이터레이터처럼 처리
class Square:
    def __getitem__(self, index):
        if index >= 5:
            raise IndexError
        return index ** 2

for n in Square():
    print(n)
# 0, 1, 4, 9, 16

두 방식의 차이를 정리하면 다음과 같습니다.