[점프 투 파이썬] 7장 - 파이썬 날아오르기: 이터레이터와 제너레이터

본 글은 '점프 투 파이썬(박응용)'을 바탕으로 공부한 내용을 정리한 글입니다.

이터레이터

iterator: 데이터를 하나씩 순서대로 꺼내올 수 있는 객체

• 즉, next() 함수를 사용하여 값을 하나씩 가져올 수 있고
• 모든 값을 가져오면 StopIteration 예외 발생
• for, next()로 그 값을 한 번 읽으면 그 값을 다시는 읽을 수 없음
• iterable(반복 가능) object ≠ iterator
  • iterable 객체는 상태 가지지 X → iteration 할 때마다 현재 상태 기억하는 새로운 iterator 객체 만듦
  • iterator 객체는 상태(== 현재 위치) 기억 → iteration 하면 멈췄던 곳에서 다시 시작
  • 예를 들어, for-loop은 __iter__ 호출해 iterator 객체 얻음 → 그 객체의 __next__ 계속 호출

e.g. list: iterable하지만 iterator는 아님

a = [1, 2, 3]
# print(next(a)) # TypeError: 'list' object is not an iterator
# iter(): 이터레이터로 만듦
ia = iter(a)
print(type(ia)) # <class 'list_iterator'>

for i in ia:
    print(i)
"""
1
2
3
"""

for i in ia:
    print(i) # 아무 것도 출력되지 않음

이터레이터 만들기

iterator class 만드는 핵심 규칙

1. __iter__ 메서드: 이터레이터 객체 자신 반환
   1. 이 메서드가 있어야, 파이썬이 해당 객체를 iterable 객체로 인식
2. __next__ 메서드: 다음 값 반환, 더 이상 값이 없으면 StopIteration 예외 발생
   1. self.index를 이용해 현재 위치의 값 가져옴
   2. 위치 증가
   3. 더 이상 값이 없으면 StopIteration 예외 발생

class StudentScoreIterator:
    def __init__(self, scores):
        self.scores = scores
        self.keys = list(self.scores.keys())
        self.index = 0 # 현재 위치 추적

    def __iter__(self):
        return self

    def __next__(self):
        try:
            result = self.scores[self.keys[self.index]]
            self.index += 1
            return result
        except:
            raise StopIteration

students = {'김민준': 85, '이서연': 92, '박지민': 78}

for score in StudentScoreIterator(students):
    print(score)

"""
85
92
78
"""

입력 데이터를 역순으로 출력하는 클래스

class Reverser:
    def __init__(self, data):
        self.data = data
        self.index = len(self.data) - 1

    def __iter__(self):
        return self

    def __next__(self):
        if self.index < 0:
            raise StopIteration
        result = self.data[self.index]
        self.index -= 1
        return result

my_string = "hello"
for char in Reverser(my_string):
    print(char)

"""
o
l
l
e
h
"""

제너레이터

제너레이터: 간단히 이터레이터를 만드는 함수

• __iter__, __next__ 없이 이터레이터 만들 수 있음

제너레이터의 핵심 특징

1. return 대신 yield 키워드 사용
2. yield 를 만나면 값을 반환, 함수 일시 정지
3. 다시 호출하면 일시정지했던 지점부터 계속 실행

def mygen():
    yield 0
    yield 1

g = mygen()
print(type(g)) # <class 'generator'>

print(next(g), next(g)) # 0 1

print(next(g)) # StopIteration

def countdown_generator(start_num):
    while start_num > 0:
        result = start_num
        start_num -= 1
        yield result

for num in countdown_generator(3):
    print(num)

"""
3
2
1
"""

제너레이터 표현식

• list comprehension 구문과 유사, but 튜플 사용

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

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

"""
아래 제너레이터 표현식은 
위 gen 제너레이터와 완전히 똑같이 기능함
"""
gen = (i * i for i in range(1, 1000))

제너레이터와 이터레이터

제너레이터, 이터레이터 서로 유사 → 간단한 건 제너레이터, 복잡한 건 이터레이터

제너레이터 활용하기

상황: 시간이 오래 걸리는 작업을 처리하는데, 모든 결과가 필요하지 않은 경우

list comprehension 사용시

import time

def long_time_work():
    print("work start")
    time.sleep(1) # 1초 지연
    return "done"

# 5번의 작업을 모두 실행해 리스트로 만듦
list_work = [long_time_work() for i in range(5)]
# 1번째 결과만 필요
print(list_work[0])

"""
work start
work start
work start
work start
work start
done
"""

• list_work = [long_time_work() for i in range(5)] 수행에 5초 소요
• 첫 번째 결과만 필요한데도 5초 기다려야 함

제너레이터 사용

import time

def long_time_work():
    print("work start")
    time.sleep(1)
    return "done"

# 제너레이터 표현식: 함수를 필요할 때만 실행
list_work = (long_time_work() for _ in range(5))
print(next(list_work))

"""
work start
done
"""

list_work = (long_time_work() for _ in range(5))

• 함수를 즉시 실행 X
• next(list_work) 호출 시 비로소 첫 번째 long_time_work() 실행

제너레이터가 유용한 경우

• 대용량 데이터 처리 시 메모리 절약
• 시간이 오래 걸리는 작업을 필요할 때만 실행
• 무한한 데이터 스트림 처리