Python 고급 문법 2: 이터레이터(iterator)와 제너레이터(generator)

환영합니다, Rolling Ress의 카루입니다.

​

이번에는 파이썬의 고급 문법 중 이터레이터와 제너레이터에 대해 알아보겠습니다. 이 두 개념은 반복문에서 많이 사용되므로 for문에 대해 정확하게 숙지하고 오셔야 아래 내용을 이해하기 편하실 겁니다.

​

Iterator(반복자, 이터레이터)

Python의 반복자(이하 이터레이터)는 C++의 반복자와 비슷한 개념으로 보아도 됩니다. 값을 한 번에 하나씩, 순서대로 꺼내주는 객체를 이터레이터라고 부릅니다.

​

• Iterable(반복 가능한 객체): 리스트(list), 튜플(tuple), 문자열(str), 딕셔너리(dict)와 같이 데이터가 여러 개 모여 있는 객체로, 이터레이터를 만들어낼 수 있습니다. 정확히는 내부에 __iter__ 메서드를 가지고 있습니다.
• Iterator(이터레이터): Iterable 객체에서 값을 하나씩 꺼내주는 도구입니다. 값을 꺼내는 __next__ 메서드를 가지고 있습니다.

numbers = [10, 20, 30]  # numbers는 Iterable(반복 가능한 객체)입니다.

for num in numbers: # 우리가 흔히 썼던 for문이지만...
    print(num)

# 실제로는 아래와 같은 과정을 거칩니다.
iterator = iter(numbers)  # 1. Iterable에서 이터레이터를 꺼냅니다.

print(next(iterator))  # 2. 첫 번째 값을 꺼냅니다. -> 10
print(next(iterator))  # 3. 두 번째 값을 꺼냅니다. -> 20
print(next(iterator))  # 4. 세 번째 값을 꺼냅니다. -> 30

# print(next(iterator)) # 5. 더 꺼낼 게 없으면 StopIteration 에러가 납니다!
# for 문은 이 에러가 발생하면 알아서 반복을 종료하는 것입니다.

 

코드를 하나씩 살펴보겠습니다. 우선 numbers는 세 개의 정수를 가진 리스트입니다. 리스트는 반복 가능한 객체(iterable)입니다. for문을 통해 출력할 수도 있지만, 이터레이터를 이용해 직접 출력해보겠습니다.

​

iter()는 iterable 객체 내부의 __iter__() 메서드를 호출하여 이터레이터를 받아옵니다. 즉, iterator 변수 안에는 이터레이터가 들어있습니다. 이제부턴 이 이터레이터를 사용하여 작업할 겁니다.

​

next() 함수는 내부적으로 매개변수의 __next__()를 호출합니다. 여기선 iterator 변수를 가지고 작업하므로, iterator.__next__()를 호출합니다. next() 함수를 호출할 때마다 앞에서부터 값을 하나씩 꺼내게 되며, 더 이상 꺼낼 값이 없으면 StopIteration 예외가 발생합니다.

class TicketMachine:    
    def __init__(self, limit):
        self.current = 1    # 현재 발급할 번호 (상태를 기억함)
        self.limit = limit  # 발급할 최대 번호

    def __iter__(self):
        # for 문 같은 곳에서 iter()를 호출하면 자기 자신을 내어줍니다.
        return self

    def __next__(self):
        # next()를 호출할 때마다 실행됩니다.
        if self.current <= self.limit:
            ticket_number = self.current
            self.current += 1  # 다음 번호를 위해 1 증가시킴 (상태 변경)
            return f"대기번호 {ticket_number}번"
        else:
            # 한도에 도달하면 표가 떨어졌다는 신호를 보냅니다.
            raise StopIteration

machine = TicketMachine(5) # 5번까지만 표를 주는 기계 생성

print("=== 수동 출력 ===")
print(next(machine)) # 대기번호 1번
print(next(machine)) # 대기번호 2번

print("=== for문 출력 ===")
for ticket in machine:
    print(ticket)    # 대기번호 3, 4, 5번

이터레이터를 직접 만들 수도 있습니다. 기본적으로 __iter__ (자기 자신을 반환), __next__ (다음 값을 꺼냄) 이렇게 두 개의 메서드를 보유해야 합니다. __iter__에는 자기 자신을 반환하는 코드만 있으면 되며, __next__는 현재 요소를 반환하고 만약 더 이상 반환할 내용이 없다면 StopIteration 예외를 발생시키면 됩니다.

from time import sleep

class Natural:    
    def __init__(self):
        self.current = 1

    def __iter__(self):
        return self

    def __next__(self):
        c = self.current
        self.current += 1
        return c

n = Natural() # 자연수의 집합
for i in n:
    print(i)
    sleep(0.1)

이런 식으로 사용하면 마치 무한집합처럼 동작하는 코드를 만들 수 있습니다. 분명히 for문을 통해 반복하고 있는데, 저는 리스트를 선언하지 않았습니다. 설령 리스트를 선언한다고 해도 크기가 무한대인 리스트를 선언할 수는 없습니다. 하지만 이터레이터를 이용하면 실시간으로 다음 값을 계산해서 1개만 반환해도 됩니다.

​

이터레이터는 메모리 효율성에서 장점이 많지만, 클래스를 선언해야 하므로 불편합니다. 그래서 제너레이터가 등장합니다.

​

Generator(제너레이터)

제너레이터는 클래스 없이 함수만으로 이터레이터를 만듭니다. __iter__, __next__같은 메서드를 직접 정의할 필요가 없습니다. yield라는 키워드만 알면 됩니다.

​

yield는 return과 같이 함수에서 값을 반환하는 키워드입니다. return이 함수를 완전히 종료하고 값을 반환한다면, yield는 함수를 일시정지한 뒤 값을 반환하며, 다시 호출되면 멈췄던 부분부터 실행합니다. 이 점이 가장 큰 차이점입니다. 만약 일시정지+값 반환이라는 부분에서 __next__가 떠오르셨다면, 감이 좋으신 겁니다.

​

앞서 작성한 번호표 기계 프로그램을 제너레이터를 이용하여 다시 써보겠습니다.

def ticket_generator(limit):
    current = 1
    while current <= limit:
        yield f"대기번호 {current}번"  # 값을 던져주고 대기하기
        
        # 다음 next()가 호출되면 여기서부터 다시 시작하여 current를 1 증가시킴
        current += 1 

machine = ticket_generator(5)

for ticket in machine:
    print(ticket)

코드가 간결해졌습니다. yield를 쓰면 이 함수는 즉시 실행되지 않고, 제너레이터 타입만 반환합니다. machine에는 제너레이터가 들어가게 됩니다. 함수를 yield 부분까지 실행하려면 next(machine)을 이용하여 다음 값을 받아오면 됩니다. 마찬가지로 반환할 값이 더 이상 없으면(뒤에 yield가 없다면) StopIteration 예외가 발생합니다.

from time import sleep

def natural():
    n = 1
    while True:
        yield n
        n += 1

n = natural() # 자연수의 집합
for i in n:
    print(i)
    sleep(0.1)

이렇게 무한집합을 간단하게 정의할 수도 있습니다. 파이썬 첫 글에서 잠깐 설명했던 내용이기도 합니다.

​

Generator Expression

# 1. 리스트 컴프리헨션 (대괄호 사용)
list_comp = [x**2 for x in range(1, 6)]
print(list_comp)       # 출력: [1, 4, 9, 16, 25]
print(type(list_comp)) # 출력: <class 'list'>


# 2. 제너레이터 표현식 (소괄호 사용)
gen_expr = (x**2 for x in range(1, 6))
print(gen_expr)        # 출력: <generator object <genexpr> at 0x...>
print(type(gen_expr))  # 출력: <class 'generator'>

list comprehension은 마치 집합의 조건제시법처럼 리스트를 생성할 수 있는 방법입니다. 여기서 대괄호를 소괄호로 바꿔주기만 하면 제너레이터 표현식이 됩니다. 위 코드는 1부터 5까지의 숫자를 제곱하는 코드입니다. list_comp의 경우 이미 값이 다 계산되어 메모리에 저장되기 때문에 그대로 출력이 됩니다. 반면 gen_expr은 next()가 호출될 때 값이 계산되는 제너레이터이기에 값이 바로 출력되지 않습니다.

import sys

# 1천만 개의 숫자를 담는 리스트와 제너레이터 생성
list_comp = [x for x in range(10000000)]
gen_expr  = (x for x in range(10000000))

# 두 객체가 차지하는 메모리 용량(Byte) 확인
print(f"리스트 메모리: {sys.getsizeof(list_comp):,} 바이트")  
print(f"제너레이터 메모리: {sys.getsizeof(gen_expr):,} 바이트")

실시간으로 값을 계산하기 때문에 메모리도 적게 차지합니다. 대용량 데이터 처리 시 메모리를 적게 쓸 수 있는 것이 제너레이터의 장점입니다.

# 1부터 1000까지의 숫자 중 짝수만 골라서 모두 더하고 싶을 때

# X 나쁜 예: 리스트를 메모리에 굳이 만들었다가 합을 구하고 버림
total_bad = sum([x for x in range(1, 1001) if x % 2 == 0])

# O 좋은 예: 제너레이터 표현식으로 값을 하나씩 뽑으면서 바로 더함
total_good = sum(x for x in range(1, 1001) if x % 2 == 0)

print(total_good) # 250500

특히 sum, max, min 등 파이썬 내장 함수와 사용할 때 진가를 발휘합니다. total_good을 보면 sum 함수의 인수로 제너레이터 표현식을 사용합니다. 원래는 sum((x for ...))과 같이 작성해야 하지만, 소괄호가 겹치면 하나를 생략할 수 있습니다. 표현식으로 생성한 데이터를 여러 번 읽을 필요가 없다면, 다시 말해 처음부터 끝까지 딱 한 번만 순서대로 읽고 버려도 된다면 제너레이터 표현식을 사용하는 것이 좋습니다. Lazy evaluation을 통해 시간을 절약할 수 있고, 메모리 사용량도 줄일 수 있기 때문입니다.