파이썬 컴프리헨션과 제너레이터 - yield 와 메모리

이 단원에서 배우는 것

10단원까지 주문 데이터를 클래스로 세우고, 잘못된 주문 한 건이 배치 전체를 죽이지 않게 막았다. 이제 데이터 양이 문제가 된다. 주문 로그가 하루 10만 줄이면 리스트에 다 담아도 괜찮지만, 한 달치 300만 줄을 한 번에 올리면 메모리가 버티지 못한다. 이번 단원에서는 같은 반복 처리를 메모리를 쓰지 않고 하는 방법과, 그 대가로 무엇을 포기하는지를 다룬다.

• 컴프리헨션이 읽기 좋아지는 경계와 오히려 나빠지는 경계를 판단할 수 있다
• 제너레이터가 값을 언제 계산하는지 설명하고, 큰 로그 파일을 일정한 메모리로 처리할 수 있다
• 제너레이터가 한 번만 소비된다는 성질에서 오는 버그를 미리 알아본다

왜 필요한가

주문 로그에서 12000원 이상인 항목의 금액만 뽑는다고 하자. 지금까지 써 온 방식은 이렇다.

amounts = []
for line in lines:
    if line.unit_price >= 12000:
        amounts.append(line.qty * line.unit_price)

네 줄 중 실제 의미가 있는 것은 조건과 계산식 두 개다. 나머지 두 줄(빈 리스트 만들기, append)은 "리스트를 만든다"는 뜻을 전달하는 데 쓰이는 사무 절차다. 컴프리헨션은 이 사무 절차를 없앤다.

amounts = [line.qty * line.unit_price for line in lines if line.unit_price >= 12000]

더 중요한 문제는 메모리다. 위 방식은 lines가 이미 전부 메모리에 있다고 전제한다. 로그 파일이 2GB라면 그 전제부터 깨진다. 제너레이터는 한 번에 한 건씩 만들어 넘겨주고 잊어버린다. 그래서 2GB 파일도 몇 MB 메모리로 처리된다. 이 단원의 두 도구는 같은 뿌리에서 나온다. 어느 쪽을 쓸지는 "결과 전체가 동시에 필요한가"로 갈린다.

문법과 예제

세 가지 컴프리헨션

from dataclasses import dataclass


@dataclass(frozen=True)
class OrderLine:
    sku: str
    qty: int
    unit_price: int

    @property
    def amount(self) -> int:
        return self.qty * self.unit_price


lines = [
    OrderLine("BOOK-01", 2, 12000),
    OrderLine("PEN-07", 5, 3000),
    OrderLine("BOOK-02", 1, 25000),
    OrderLine("PEN-07", 2, 3000),
]

# 리스트: 순서 있는 결과
amounts = [line.amount for line in lines if line.unit_price >= 12000]
print(amounts)          # [24000, 25000]

# 집합: 중복 제거
skus = {line.sku for line in lines}
print(sorted(skus))     # ['BOOK-01', 'BOOK-02', 'PEN-07']

# 딕셔너리: 키로 찾기 위한 표
price_of = {line.sku: line.unit_price for line in lines}
print(price_of)         # {'BOOK-01': 12000, 'PEN-07': 3000, 'BOOK-02': 25000}

딕셔너리 컴프리헨션에서 키가 겹치면 나중 값이 이긴다. 위에서 PEN-07이 두 번 나오지만 결과에는 하나만 남는다. 단가가 같아서 문제가 안 됐을 뿐, 값이 달랐다면 조용히 하나가 사라진다. 겹칠 수 있는 키로 표를 만들 때는 13단원에서 다룰 defaultdict가 맞는 도구다.

읽기 좋은 경계

컴프리헨션은 짧아서 좋은 게 아니라 한 줄이 한 문장으로 읽혀서 좋다. 그 성질이 깨지는 순간 장점이 사라진다.

# 읽힌다: "각 항목의 금액을, 단가 12000 이상인 것만"
[line.amount for line in lines if line.unit_price >= 12000]

# 안 읽힌다: 조건이 세 개, 표현식에 삼항 연산자까지
[line.amount * 0.9 if line.qty >= 3 else line.amount
 for line in lines
 if line.unit_price >= 12000 and line.sku.startswith("BOOK") and line.qty > 0]

두 번째는 for 문으로 푸는 편이 낫다. 판단 기준을 하나 잡자면 중첩 반복이 들어가거나, 조건이 두 개를 넘거나, 한 줄에 안 들어가면 되돌린다. 특히 중첩은 순서가 헷갈린다.

orders = [["BOOK-01", "PEN-07"], ["BOOK-02"]]

# for 를 쓴 순서와 같은 순서로 왼쪽부터 읽는다
flat = [sku for order in orders for sku in order]
print(flat)   # ['BOOK-01', 'PEN-07', 'BOOK-02']

for order in orders가 바깥, for sku in order가 안쪽이다. 바깥 루프가 먼저 오는데, 결과 표현식(sku)은 맨 앞에 있다. 이 배치 때문에 두 개를 넘어가면 사람이 못 읽는다. 13단원의 itertools.chain.from_iterable이 같은 일을 더 분명하게 한다.

제너레이터 표현식: 대괄호를 소괄호로

import sys

# 리스트: 100만 개를 전부 메모리에 만든다
listcomp = [n * 2 for n in range(1_000_000)]
print(sys.getsizeof(listcomp))   # 8448728 (약 8MB)

# 제너레이터: 만들지 않는다. 만드는 방법만 들고 있다
genexp = (n * 2 for n in range(1_000_000))
print(sys.getsizeof(genexp))     # 200 정도

print(sum(genexp))               # 999999000000
print(sum(genexp))               # 0  <- 이미 다 썼다

마지막 두 줄이 이 단원에서 가장 중요한 지점이다. 제너레이터는 한 번 순회하면 끝난다. 되감기가 없다. 두 번째 sum은 에러도 없이 0을 돌려준다. 조용히 틀린 답이 나오는 부류의 버그다.

함수 인자로 넘길 때는 괄호를 생략할 수 있다.

lines = [OrderLine("BOOK-01", 2, 12000), OrderLine("PEN-07", 5, 3000)]
total = sum(line.amount for line in lines)     # 괄호 한 겹 생략
print(total)   # 39000

yield: 함수를 멈췄다 이어서 실행하기

yield가 하나라도 있으면 그 함수는 호출해도 몸통이 실행되지 않는다. 제너레이터 객체만 돌려주고, 실제 실행은 값을 하나 요청할 때마다 yield까지 진행하고 그 자리에서 멈춘다.

def counter(n):
    print("시작")
    for i in range(n):
        print(f"  yield 직전 {i}")
        yield i
    print("끝")


gen = counter(3)
print("아직 아무것도 실행되지 않았다")
print(next(gen))
print(next(gen))

출력은 이렇게 된다.

아직 아무것도 실행되지 않았다
시작
  yield 직전 0
0
  yield 직전 1
1

"시작"이 counter(3) 호출 때가 아니라 첫 next() 때 찍힌다. 이 지연 실행이 제너레이터의 전부다.

실무형: 로그 파일을 파이프라인으로 처리하기

주문 로그를 읽어 집계하는 코드를 제너레이터 단계로 나눈다. 각 단계는 자기 일만 하고, 전체가 한 줄씩 흘러간다.

from pathlib import Path


def read_lines(path: Path):
    """파일을 한 줄씩 내보낸다. 파일 전체를 메모리에 올리지 않는다."""
    with path.open(encoding="utf-8") as f:
        for raw in f:
            yield raw.rstrip("\n")


def parse_orders(raw_lines):
    """'ORD-1001,BOOK-01,2,12000' 형식을 OrderLine 으로 바꾼다."""
    for raw in raw_lines:
        if not raw or raw.startswith("#"):
            continue
        order_id, sku, qty, price = raw.split(",")
        yield order_id, OrderLine(sku, int(qty), int(price))


def only_books(pairs):
    for order_id, line in pairs:
        if line.sku.startswith("BOOK"):
            yield order_id, line


path = Path("orders.csv")
path.write_text(
    "# order_id,sku,qty,unit_price\n"
    "ORD-1001,BOOK-01,2,12000\n"
    "ORD-1001,PEN-07,5,3000\n"
    "ORD-1002,BOOK-02,1,25000\n",
    encoding="utf-8",
)

pipeline = only_books(parse_orders(read_lines(path)))
print(sum(line.amount for _, line in pipeline))   # 49000

pipeline을 만드는 시점에는 파일이 아직 열리지도 않았다. 마지막 sum이 값을 당겨오기 시작할 때 비로소 파일이 열리고, 한 줄씩 세 단계를 통과한다. 파일이 100GB여도 메모리 사용량은 같다. 단계를 추가하거나 순서를 바꾸는 것도 함수 하나를 끼워 넣으면 된다.

yield from

제너레이터 안에서 다른 제너레이터의 값을 전부 흘려보낼 때 쓴다.

def read_all(paths):
    for path in paths:
        yield from read_lines(path)      # for 로 풀어 쓴 것과 같다

for x in read_lines(path): yield x와 결과가 같지만, yield from은 send()와 예외까지 안쪽 제너레이터에 그대로 전달한다. 일반적인 반복 처리에서는 짧아서 쓰고, 코루틴을 다룰 때는 의미가 달라진다.

어느 쪽을 쓸지 정하는 기준

상황	선택	이유
결과를 두 번 이상 순회한다	리스트	제너레이터는 두 번째부터 비어 있다
len(), 인덱싱, 정렬이 필요하다	리스트	제너레이터는 셋 다 지원하지 않는다
sum, max, any에 바로 넣는다	제너레이터	중간 리스트를 만들 이유가 없다
입력이 파일, 네트워크, DB 커서다	제너레이터	전체 크기를 모르거나 메모리를 넘는다
앞의 몇 건만 보고 끝낼 수도 있다	제너레이터	나머지는 계산조차 하지 않는다

제너레이터를 함수형 스타일로 더 밀고 나가는 방법은 함수형 프로그래밍 HOWTO에 정리돼 있다. 13단원에서 다룰 itertools가 그 연장선이다.

실무에서 자주 틀리는 것

1. 제너레이터를 두 번 순회한다

def load_lines():
    yield OrderLine("BOOK-01", 2, 12000)
    yield OrderLine("PEN-07", 5, 3000)


lines = load_lines()
total = sum(line.amount for line in lines)
count = sum(1 for line in lines)          # 이미 소비됐다

print(total)              # 39000
print(count)              # 0
print(total / count)      # ZeroDivisionError

평균 계산 같은 자리에서 ZeroDivisionError로 터지면 그나마 다행이고, 나쁜 경우는 개수가 0이어도 코드가 넘어가서 통계가 조용히 0으로 찍힌다. 두 번 봐야 한다면 lines = list(load_lines())로 한 번 받아 두거나, 함수를 두 번 호출해 새 제너레이터를 만든다. 제너레이터를 반환하는 함수는 이름에 그 성질이 드러나게 짓는 편이 사고를 줄인다.

2. with 블록 밖에서 제너레이터를 소비한다

def bad_read(path):
    with open(path, encoding="utf-8") as f:
        return (line.rstrip("\n") for line in f)   # return 이다, yield 가 아니다


gen = bad_read("orders.csv")
print(list(gen))   # ValueError: I/O operation on closed file.

return하는 순간 with 블록이 끝나면서 파일이 닫힌다. 제너레이터는 아직 한 줄도 안 읽었는데 소스가 사라진 것이다. 앞의 read_lines처럼 with 안에서 yield 해야 한다. 그러면 파일은 제너레이터가 다 소비되거나 가비지 컬렉션될 때까지 열려 있다.

3. 컴프리헨션을 부작용 목적으로 쓴다

[save_to_db(line) for line in lines]     # 결과 리스트를 아무도 안 쓴다

None이 10만 개 든 리스트를 만들고 버린다. 메모리도 아깝지만 읽는 사람이 "이 리스트가 어디 쓰이나" 찾느라 시간을 쓴다. 저장이 목적이면 for line in lines: save_to_db(line)이라고 쓴다. 컴프리헨션은 값을 만들 때만 쓴다.

4. any/all 에 리스트 컴프리헨션을 넣는다

if any([is_invalid(line) for line in lines]):   # 전부 검사한다
    ...

if any(is_invalid(line) for line in lines):     # 첫 True 에서 멈춘다
    ...

any와 all은 답이 확정되는 순간 멈추는데, 대괄호를 쓰면 그전에 리스트를 완성해 버려서 그 장점이 통째로 사라진다. is_invalid가 DB를 조회한다면 첫 줄에서 끝날 일을 10만 번 한다. 대괄호 하나 차이다.

스스로 확인하기

1. 아래 코드는 합계와 최대를 한 줄에 찍으려 했지만 예외로 죽는다. 무엇이 문제이며 어떻게 고치는가?

   amounts = (line.amount for line in lines)
   print(f"합계 {sum(amounts)}, 최대 {max(amounts)}")
   

2. 2GB짜리 주문 로그에서 BOOK으로 시작하는 SKU가 하나라도 있는지만 알고 싶다. 리스트 컴프리헨션 대신 무엇을 쓰고, 그 차이가 실행 시간에 어떻게 나타나는가?
3. {line.sku: line.unit_price for line in lines}로 만든 표에서 어떤 SKU의 단가가 실제 데이터와 다르게 나올 수 있는 경우는?

정답

1. sum(amounts)가 제너레이터를 전부 소비하므로 max(amounts)는 빈 것을 받아 "빈 대상에서 최댓값을 구할 수 없다"는 ValueError로 죽는다(정확한 문구는 파이썬 버전마다 다르다). 두 번 순회해야 하므로 amounts = [line.amount for line in lines]로 리스트에 담는다.
2. any(line.sku.startswith("BOOK") for line in parse_orders(read_lines(path)))를 쓴다. 조건을 만족하는 첫 줄에서 즉시 True를 반환하고 파일 읽기를 멈춘다. 리스트 컴프리헨션은 답이 첫 줄에서 정해져도 2GB를 끝까지 읽는다.
3. 같은 SKU가 여러 항목에 나오면서 단가가 다를 때다. 딕셔너리 컴프리헨션은 뒤에 나온 값으로 덮어쓰므로 마지막 항목의 단가만 남고 앞의 값은 경고 없이 사라진다.