Profile picture

[Python] 제너레이터(Generator)와 yield 키워드

JaehyoJJAng2024년 06월 04일

Yield

일반적으로 프로그래밍에서의 함수는 어떤 결과 값을 return 키워드를 이용해 값을 반환 한다.

하지만 파이썬에서는 함수에서 yield 키워드를 사용하면 다른 방식으로 값을 반환할 수도 있다.


간단한 예제 코드를 보며 이해해보자.

먼저, 문자열 q,w,e를 결과 값으로 반환하는 함수를 작성해보자.

def return_str() -> list[str]:
    return list('qwe')

이번에는 위 함수를 yield 키워드를 이용해 작성해보자.

def yield_str():
    yield 'q'
    yield 'w'
    yield 'e'

두 함수의 차이는 return 키워드를 사용할 때는 결과 값을 딱 한 번만 제공하는데,

yield 키워드는 여러번 나누어 결과 값을 제공하고 있다.


for 루프를 사용하여 위 함수를 호출해 얻은 결과를 출력해보자.

for rst in return_str():
    print(rst)
q
w
e

for ch in yield_str():
    print(ch)
q
w
e

함수를 사용하는 측면에서 보면 두 함수는 큰 차이가 없어보인다.


함수를 호출한 결과 값을 바로 출력해서 각 함수가 정확히 뭘 리턴하는지 알아보자.

print(return_str())
['q', 'w', 'e']

print(yield_str())
<generator object yield_str at 0x0000022295624B40>

return_str() 함수는 리스트를 반환하고 있고, yield_str() 함수는 제너레이터를 반환하고 있다.


여기서 yield 키워드를 사용하면 제너레이터를 반환한다는 것을 알 수 있는데, 제너레이터는 어떤 개념인걸까?


제너레이터

파이썬에서의 제너레이터는 여러 개의 데이터를 미리 만들어 놓지 않고 필요할 때마다 하나씩 만들어낼 수 있는 객체를 의미한다.

위에서 작성한 예제 코드를 참고해 알파벳 하나를 만드는데 1초가 걸리게끔 코드를 수정해보자.

import time

def return_str() -> list[str]:
    alph = []
    for word in "qwe":
        time.sleep(1)
        alph.append(word)
    return alph

위 함수를 호출한 결과를 for 루프로 돌려보면 3초가 흐른 후에 q, w, e가 한 번에 출력된다.

for rst in return_str():
    print(rst)
# 3초 경과
q
w
e

이번에는 yield 키워드를 사용해 동일한 결과 값을 제공하는 함수를 작성해보자.

import time

def yield_str():
    for word in "qwe":
        time.sleep(1)
        yield word

위 함수를 호출한 결과를 for 루프로 돌려보면 1초 후에 q를 출력하고, 또 1초 후에 w를 출력하고, 또 1초 후에 e가 출력된다.

for rst in yield_str():
    print(rst)
# 1초 경과
q
# 2초 경과
w
# 3초 경과
e

만약에 세개의 문자열이 아닌 백개, 그 이상의 문자열을 제공해야 하는 경우에는 어떨까?

yield 키워드를 사용하지 않은 함수에서는 첫 번째 결과 값을 얻는데 백초, 그 이상의 시간이 걸리는 반면에,

yield 키워드를 사용한 함수에서는 항상 일초가 걸릴 것이다.

즉, 제너레이터를 통해서 결과 값을 나누어 얻을 수 있기 때문에 성능 측면에서 이점이 존재한다.


데이터 무한으로 생성해보기

제너레이터를 사용하면 이론적으로 무한한 데이터를 계속해서 만들어볼 수도 있다.

예를 들어, Q, W, E를 계속해서 무한하게 출력하는 함수를 작성해보자.

def yield_infinite_qwe():
    while True:
        yield 'Q'
        yield 'W'
        yield 'E'

이 함수를 호출한 결과를 for 루프로 돌리면 Q, W, E가 화면에 계속 출력될 것이다.

for rst in yield_infinite_qwe():
    print(rst)

이렇게 무한한 데이터를 제공하는 함수를 yield 키워드 없이 작성하는 것은 불가능에 가깝다.

컴퓨터의 물리적 메모리에는 한계가 있고, 아무리 큰 리스트를 생성하더라도 이 한계를 초과할 수는 없기 때문이다.


yield from

제너레이터를 반환하는 함수를 작성하다보면 아래와 같이 리스트를 제너레이터로 변환해야 할 일이 생길 수 있다.

def yield_qwe():
    for ch in ["q", "w", "e"]:
        yield ch

이 때 yield from를 사용하면 리스트를 바로 제너레이터로 변환할 수 있어 매우 편리하다.

def yield_qwe():
    yield from ["q", "w", "e"]

예제 코드

지금까지 파이썬 제너레이터가 사용되는 간단한 예제들을 다뤄봤는데,

이번에는 실제로 사용되는 케이스들에 대한 예시를 몇 가지 살펴보도록 하자.


데이터 스트리밍 및 파일 처리

제너레이터의 가장 큰 특징인, 한번에 하나의 출력만 수행한다는 점을 활용해,

큰 데이터셋을 처리할 때 전체 데이터를 한 번에 메모리에 로드하는 대신, 필요한 부분만 순차적으로 처리할 수 있다.

def read_file(file_name: str):
    with open(file_name, 'r', encoding='utf-8') as fp:
        for line in fp.readlines():
            yield line

# 파일에서 한 줄씩 읽어오기
for line in read_file(file_name='test.txt'):
    print(line)

위 코드에서의 read_file 함수에서는 로드한 파일에 대해 yield 키워드를 사용해 한 줄씩 출력한다.

덕분에 한 번에 데이터를 로드하지 않고도, 필요할 때마다 한 줄씩 꺼내서 출력이 가능하다.


비슷한 예시를 한가지 더 들어보자.

def read_file_in_chunks(file_name: str, chunk_size: int=1024):
    with open(file_name, 'r', encoding='utf-8') as fp:
        while True:
            chunk = fp.read(chunk_size)
            if not chunk:
                break
            yield chunk

# 파일을 청크 단위로 읽어오기
for chunk in read_file_in_chunks(file_name='test.txt'):
    print(chunk)

파일을 chunk 단위로 읽어 필요한 만큼 반복적으로 불러올 수 있다.

예를 들어, 대용량의 로그 파일을 읽어와 특정 키워드를 체크한다고 할 때,

전체 파일에 대한 내용을 읽어와 분석한다면 메모리가 크게 사용되어 비효율적일 것이다.

이 때 yield 키워드를 사용하면 한 번에 하나의 line을 불러와 해당 라인에 대한 분석이 가능하기 때문에 메모리를 최적화 하는데 매우 유용할 것으로 보인다.


Loading script...