파이썬 – 제너레이터 (Generator)

9 분 소요

앞서 다룬 파이썬 – 데코레이터 (Decorator)에 이어 이번 글에서는 파이썬의 제너레이터(Generator)를 정리합니다.

파이썬은 비교적 배우기 쉬운 언어이지만, 파이썬 초보자가 공통적으로 어려워하는 개념이 몇 가지 있습니다. 그중 하나가 제너레이터와 yield입니다.

제너레이터의 사전적 의미는 “발전기” 또는 “무언가를 만들어내는 대상"입니다. 위키피디아의 컴퓨터 과학적 정의는 다음과 같습니다.

제너레이터는 반복자(iterator)와 같은 루프의 작용을 컨트롤하기 위해 쓰여지는 특별한 함수 또는 루틴이다. 사실 모든 제너레이터는 반복자이다. 제너레이터는 배열이나 리스트를 리턴하는 함수와 비슷하며, 호출을 할 수 있는 파라메터를 가지고 있고, 연속적인 값들을 만들어 낸다. 하지만 한번에 모든 값을 포함한 배열을 만들어서 리턴하는 대신에 yield 구문을 이용해 한 번 호출될 때마다 하나의 값만을 리턴하고, 이런 이유로 일반 반복자에 비해 아주 작은 메모리를 필요로 한다. 간단히 얘기하면 제너레이터는 반복자와 같은 역할을 하는 함수이다.

일반 함수는 호출되면 첫 번째 행부터 실행되어 return 구문이나 예외(exception), 혹은 (리턴이 없는 함수라면) 마지막 구문에 도달할 때까지 실행된 뒤 호출자(caller)에게 제어권을 돌려줍니다. 그리고 그 시점에 함수가 보유하던 내부 함수나 로컬 변수는 모두 메모리에서 사라지며, 같은 함수를 다시 호출하면 처음부터 새로 시작합니다.

그러나 한 번에 모든 일을 처리한 뒤 사라지는 함수가 아니라, 한 단계의 작업을 마친 뒤 자신의 상태를 기억한 채 대기하다가 다시 호출되면 이어서 처리하는 함수가 필요해졌습니다. 이를 위해 만들어진 것이 제너레이터입니다. 제너레이터를 사용하면 일반 함수보다 메모리 사용을 크게 줄일 수 있고, 경우에 따라 더 나은 퍼포먼스도 얻을 수 있습니다. 예제로 제너레이터의 동작을 살펴봅니다.

작업 디렉터리에 generator.py 파일을 만들고 다음 코드를 저장합니다.

generator.py
def square_numbers(nums):
result = []
for i in nums:
result.append(i * i)
return result

my_nums = square_numbers([1, 2, 3, 4, 5])

print(my_nums)

간단한 함수를 정의해 호출하는 코드입니다. 인자로 받은 리스트를 for 루프로 순회하면서 i * i의 결과로 새로운 리스트를 만들어 리턴합니다.

터미널에서 generator.py가 저장된 디렉터리로 이동한 뒤 다음 명령어로 프로그램을 실행합니다.

$ python generator.py
[1, 4, 9, 16, 25]

새 리스트가 결과로 리턴되었습니다.

이 코드를 제너레이터로 변환합니다.

generator.py
def square_numbers(nums):
for i in nums:
yield i * i

my_nums = square_numbers([1, 2, 3, 4, 5])  #1

print(my_nums)
$ python generator.py
<generator object square_numbers at 0x0000016B17E19EB0>

제너레이터 오브젝트가 리턴되었습니다. 제너레이터는 리턴할 모든 값을 메모리에 보관하지 않기 때문에 일반 함수처럼 리스트 형태로 한 번에 보이지 않습니다. 제너레이터는 호출될 때마다 하나의 값만 전달(yield)합니다. 즉, 위 #1 시점에서는 아직 아무 계산도 수행되지 않은 채 다음 값을 요청받기를 기다리는 상태입니다. 직접 확인해 봅니다.

generator.py
def square_numbers(nums):
for i in nums:
yield i * i

my_nums = square_numbers([1, 2, 3, 4, 5])

print(next(my_nums))
$ python generator.py
1

next() 함수로 다음 값을 요청했더니 1이 반환되었습니다. 이어서 여러 번 더 호출해 봅니다.

generator.py
def square_numbers(nums):
for i in nums:
yield i * i

my_nums = square_numbers([1, 2, 3, 4, 5])

print(next(my_nums))
print(next(my_nums))
print(next(my_nums))
print(next(my_nums))
print(next(my_nums))
$ python generator.py
1
4
9
16
25

앞서 일반 함수가 리턴했던 리스트의 값이 모두 출력되었습니다. 여기서 next() 함수를 한 번 더 호출하면 어떻게 될까요?

generator.py
def square_numbers(nums):
for i in nums:
yield i * i

my_nums = square_numbers([1, 2, 3, 4, 5])

print(next(my_nums))
print(next(my_nums))
print(next(my_nums))
print(next(my_nums))
print(next(my_nums))
print(next(my_nums))
$ python generator.py
1
4
9
16
25
Traceback (most recent call last):
File "generator.py", line 12, in <module>
print(next(my_nums))
StopIteration

StopIteration 예외가 발생하였습니다. 더 이상 전달할 값이 없다는 뜻입니다.

제너레이터는 일반적으로 for 루프를 통해 사용합니다. 다음 예제로 살펴봅니다.

generator.py
def square_numbers(nums):
for i in nums:
yield i * i

my_nums = square_numbers([1, 2, 3, 4, 5])

for num in my_nums:
print(num)
$ python generator.py
1
4
9
16
25

이번에는 모든 값이 출력되고 StopIteration 예외도 발생하지 않았습니다. for 루프는 어디에서 멈춰야 하는지 자체적으로 판단하기 때문입니다.

제너레이터가 일반 함수에 비해 가지는 장점 중 하나는 코드가 단순해진다는 점입니다. 파이썬 철학을 정리한 “The Zen of Python"의 세 번째 항목에는 “복잡한 것보다 단순한 것이 더 낫다"라고 적혀 있습니다.

list comprehension을 사용하면 위 코드를 한층 더 간결하게 작성할 수 있습니다.

generator.py
my_nums = [x*x for x in [1, 2, 3, 4, 5]]

print(my_nums)

for num in my_nums:
print(num)
$ python generator.py
[1, 4, 9, 16, 25]
1
4
9
16
25

첫 번째 예제의 일반 함수와 동일한 리스트를 리턴합니다. 같은 구문에서 약간만 바꾸면 제너레이터를 만들 수 있습니다.

generator.py
my_nums = (x*x for x in [1, 2, 3, 4, 5])  #1

print(my_nums)

for num in my_nums:
print(num)
$ python generator.py
<generator object <genexpr> at 0x1007c8f50>
1
4
9
16
25

#1에서 []()로 바꾸자 제너레이터가 생성되었습니다.

for 루프를 사용하지 않고 제너레이터의 데이터 전체를 한 번에 확인하려면 제너레이터를 리스트로 변환하면 됩니다.

generator.py
my_nums = (x*x for x in [1, 2, 3, 4, 5])  # 제너레이터 생성

print(my_nums)
print(list(my_nums))  # 제너레이터를 리스트로 변형
$ python generator.py
<generator object <genexpr> at 0x0000026FD7A99EB0>
[1, 4, 9, 16, 25]

정상적으로 리스트로 변환되어 출력되었습니다. 다만 한 가지 주의할 점이 있습니다. 제너레이터를 리스트로 변환하는 순간 제너레이터가 가진 장점을 모두 잃게 됩니다. 가장 핵심적인 장점은 메모리 효율입니다. 앞서 설명했듯 제너레이터는 모든 결과값을 메모리에 저장하지 않기 때문에 메모리 사용량이 적습니다. 예제로 확인해 봅니다.

generator.py
from __future__ import division
import os
import psutil
import random
import time

names = ['최용호', '지길정', '진영욱', '김세훈', '오세훈', '김민우']
majors = ['컴퓨터 공학', '국문학', '영문학', '수학', '정치']

process = psutil.Process(os.getpid())
mem_before = process.memory_info().rss / 1024 / 1024


def people_list(num_people):
result = []
for i in range(num_people):
person = {
'id': i,
'name': random.choice(names),
'major': random.choice(majors)
}
result.append(person)
return result


def people_generator(num_people):
for i in range(num_people):
person = {
'id': i,
'name': random.choice(names),
'major': random.choice(majors)
}
yield person

t1 = time.time()
people = people_list(1000000)  # 1 people_list를 호출
t2 = time.time()
mem_after = process.memory_info().rss / 1024 / 1024
total_time = t2 - t1

print('시작 전 메모리 사용량: {} MB'.format(mem_before))
print('종료 후 메모리 사용량: {} MB'.format(mem_after))
print('총 소요된 시간: {:.6f} 초'.format(total_time))
$ python generator.py
시작 전 메모리 사용량: 13.76171875 MB
종료 후 메모리 사용량: 284.30078125 MB
총 소요된 시간: 1.215000 초

먼저 #1에서 people_list(1000000)를 호출해 백만 명의 학생 정보가 담긴 리스트를 생성했습니다. 메모리 사용량이 13 MB에서 284 MB로 증가했고, 소요 시간은 1.2초였습니다. 이번에는 #1을 people_generator(1000000)로 변경해 제너레이터의 퍼포먼스를 측정해 봅니다.

generator.py
from __future__ import division
import os
import psutil
import random
import time

names = ['최용호', '지길정', '진영욱', '김세훈', '오세훈', '김민우']
majors = ['컴퓨터 공학', '국문학', '영문학', '수학', '정치']

process = psutil.Process(os.getpid())
mem_before = process.memory_info().rss / 1024 / 1024


def people_list(num_people):
result = []
for i in range(num_people):
person = {
'id': i,
'name': random.choice(names),
'major': random.choice(majors)
}
result.append(person)
return result


def people_generator(num_people):
for i in range(num_people):
person = {
'id': i,
'name': random.choice(names),
'major': random.choice(majors)
}
yield person

t1 = time.time()
people = people_generator(1000000)  # 1 people_generator를 호출
t2 = time.time()
mem_after = process.memory_info().rss / 1024 / 1024
total_time = t2 - t1

print('시작 전 메모리 사용량: {} MB'.format(mem_before))
print('종료 후 메모리 사용량: {} MB'.format(mem_after))
print('총 소요된 시간: {:.6f} 초'.format(total_time))
$ python generator.py
시작 전 메모리 사용량: 13.75390625 MB
종료 후 메모리 사용량: 13.7578125 MB
총 소요된 시간: 0.000000 초

메모리 사용량의 변화는 없었고 소요 시간도 0.1초 미만이었습니다. 제너레이터를 사용하면 메모리 사용량이 적을 뿐 아니라 제너레이터 오브젝트 생성 자체가 리스트 오브젝트 생성보다 빠르다는 점이 확인되었습니다.

다만 생성된 오브젝트를 실제로 사용해 데이터를 처리할 때는 결과가 달라질 수 있습니다. 먼저 생성된 리스트를 사용한 for 루프를 측정해 봅니다.

generator.py
from __future__ import division
import os
import psutil
import random
import time

names = ['최용호', '지길정', '진영욱', '김세훈', '오세훈', '김민우']
majors = ['컴퓨터 공학', '국문학', '영문학', '수학', '정치']

process = psutil.Process(os.getpid())
mem_before = process.memory_info().rss / 1024 / 1024


def people_list(num_people):
result = []
for i in range(num_people):
person = {
'id': i,
'name': random.choice(names),
'major': random.choice(majors)
}
result.append(person)
return result


def people_generator(num_people):
for i in range(num_people):
person = {
'id': i,
'name': random.choice(names),
'major': random.choice(majors)
}
yield person

t1 = time.time()

people = people_list(1000000)

# 리스트를 사용하여 for loop 실행
for p in people:
print(p)

t2 = time.time()
mem_after = process.memory_info().rss / 1024 / 1024
total_time = t2 - t1

print('시작 전 메모리 사용량: {} MB'.format(mem_before))
print('종료 후 메모리 사용량: {} MB'.format(mem_after))
print('총 소요된 시간: {:.6f} 초'.format(total_time))
$ python generator.py
{'id': 999998, 'name': '진영욱', 'major': '영문학'}
{'id': 999999, 'name': '진영욱', 'major': '컴퓨터 공학'}
{'id': 999999, 'name': '진영욱', 'major': '컴퓨터 공학'}
시작 전 메모리 사용량: 13.7578125 MB
종료 후 메모리 사용량: 285.84765625 MB
총 소요된 시간: 97.907999 초

메모리 사용량은 13 MB에서 284 MB로 증가했고 소요 시간은 97.9초였습니다.

이번에는 생성된 제너레이터를 사용한 for 루프를 측정합니다.

generator.py
from __future__ import division
import os
import psutil
import random
import time

names = ['최용호', '지길정', '진영욱', '김세훈', '오세훈', '김민우']
majors = ['컴퓨터 공학', '국문학', '영문학', '수학', '정치']

process = psutil.Process(os.getpid())
mem_before = process.memory_info().rss / 1024 / 1024


def people_list(num_people):
result = []
for i in range(num_people):
person = {
'id': i,
'name': random.choice(names),
'major': random.choice(majors)
}
result.append(person)
return result


def people_generator(num_people):
for i in range(num_people):
person = {
'id': i,
'name': random.choice(names),
'major': random.choice(majors)
}
yield person

t1 = time.time()

people = people_generator(1000000)  # 1 people_generator를 호출

# 제너레이터를 사용하여 for loop 실행
for p in people:
print(p)

t2 = time.time()
mem_after = process.memory_info().rss / 1024 / 1024
total_time = t2 - t1

print('시작 전 메모리 사용량: {} MB'.format(mem_before))
print('종료 후 메모리 사용량: {} MB'.format(mem_after))
print('총 소요된 시간: {:.6f} 초'.format(total_time))
$ python generator.py
{'id': 999997, 'name': '진영욱', 'major': '컴퓨터 공학'}
{'id': 999998, 'name': '오세훈', 'major': '컴퓨터 공학'}
{'id': 999999, 'name': '지길정', 'major': '영문학'}
시작 전 메모리 사용량: 13.76171875 MB
종료 후 메모리 사용량: 13.75390625 MB
총 소요된 시간: 102.774121 초

메모리 소비는 발생하지 않았지만 소요 시간은 102.7초로, 리스트 사용 시보다 약 5초 느린 결과가 나왔습니다.

정리하면, 실행 시간보다 메모리 소비를 줄여야 하는 경우에는 제너레이터를, 메모리보다 실행 시간이 더 중요한 경우에는 리스트를 사용하는 것이 적절합니다.

다만 본 예제보다 훨씬 큰 데이터를 병렬로 처리해야 하는 상황에서는 약간의 시간을 단축하는 것보다 제한된 리소스를 효율적으로 사용하는 쪽이 더 중요한 경우가 많습니다.

다음 글에서는 객체지향 프로그래밍(OOP, Object Oriented Programming)을 다룹니다.

X