파이썬 제너레이터와 yield로 대용량 파일을 메모리 효율적으로 읽기

왜 이 내용을 정리하나

파이썬으로 로그 파일이나 CSV 파일을 처리하다 보면, 파일 크기가 수 GB를 넘는 경우를 자주 만난다. 이런 파일을 read()로 통째로 읽으면 전체 내용이 메모리에 올라가서 프로그램이 느려지거나 메모리 부족 오류가 날 수 있다. 그래서 파일을 한 줄씩 읽으면서 필요한 만큼만 메모리에 두는 방법이 필요하다. 이 노트에서는 제너레이터와 yield가 이 문제에 어떻게 쓰이는지 정리해 본다.

개념 설명

제너레이터(generator)는 값을 한꺼번에 만들어 리스트에 담아 두지 않고, 요청될 때마다 하나씩 만들어 내는 이터레이터이다. 함수 안에서 yield 키워드를 쓰면 그 함수는 일반 함수가 아니라 제너레이터 함수가 된다. yield는 값을 호출한 쪽에 돌려주고 함수의 실행 상태를 그대로 멈춰 두며, 다음 값을 요구받으면 멈춘 자리에서 이어서 실행한다. 즉 return과 달리 함수가 끝나지 않는다.

파일 객체 자체도 이터레이터라서 for line in f:처럼 쓰면 한 번에 한 줄씩 읽는다. 여기에 yield를 결합하면 줄을 읽고 가공한 결과를 하나씩 흘려보내는 파이프라인을 만들 수 있다.

코드 예제

아래 코드는 큰 텍스트 파일에서 특정 단어가 포함된 줄만 골라내는 제너레이터이다.

def read_lines(path):
    with open(path, encoding="utf-8") as f:
        for line in f:
            yield line.rstrip("\n")

def filter_error(lines):
    for line in lines:
        if "ERROR" in line:
            yield line

if __name__ == "__main__":
    count = 0
    for entry in filter_error(read_lines("big.log")):
        count += 1
        if count <= 3:
            print(entry)
    print("총 개수:", count)

이 코드를 실행하면 big.log에서 ERROR가 들어간 줄 중 앞의 세 줄이 출력되고, 마지막에 전체 개수가 출력된다. read_lines는 호출 시점에 파일을 열지 않고, for문이 값을 요구할 때마다 한 줄씩 읽는다. 따라서 파일 크기와 상관없이 메모리에는 현재 줄 하나 정도만 올라간다. filter_error도 같은 방식으로 동작하여 조건에 맞는 줄만 다음 단계로 넘긴다.

헷갈리는 점

• 제너레이터는 한 번 소진되면 다시 순회할 수 없다. 같은 결과를 두 번 쓰려면 다시 호출해야 한다.
• yield가 있는 함수를 그냥 호출하면 실행되지 않고 제너레이터 객체만 반환된다. 값을 얻으려면 for나 next()로 순회해야 한다.
• list(read_lines(...))처럼 전부 리스트로 바꾸면 메모리 이점이 사라진다.
• 파일 핸들은 제너레이터가 끝까지 소비되기 전까지 열려 있다. 중간에 멈출 때는 with 블록의 생명주기를 함께 고려해야 한다.

정리

제너레이터와 yield를 쓰면 대용량 파일을 한 줄씩 읽어 필요한 부분만 처리할 수 있다. 여러 단계의 제너레이터를 연결하면 읽기, 필터링, 변환을 메모리 부담 없이 조합할 수 있다. 다만 전체 데이터를 여러 번 써야 하는 경우에는 구조를 다시 생각해야 한다.

#파이썬 #제너레이터 #yield #대용량파일 #메모리효율