Python의 List, Dict, Set Comprehension: 간결하면서도 읽기 쉽게

1. 개요

Comprehension은 Python의 특징적인 문법입니다. 이를 통해 하나의 줄 안에서 iterable로부터 list, dict 또는 set을 생성할 수 있습니다. 올바르게 작성하면 코드가 for 루프보다 짧고 명확합니다. 남용하면 코드가 긴 한 줄이 되어 이해하기 어렵습니다.

이 글에서는 list, dict, set comprehension과 generator expression을 구체적인 예제로 살펴봅니다. 이 글을 다 읽고 나면, comprehension을 언제 사용해야 하는지, 언제 for 루프로 돌아가야 하는지, 그리고 데이터 변환 로직을 어떻게 테스트할지 알게 될 것입니다. 모든 예제는 Python 3.12에서 실행됩니다.

2. List comprehension

2.1. 기본 문법

주문 목록이 있고 각 주문의 총액을 가져와야 한다고 가정해봅시다. 전통적인 for 루프를 사용하는 방법은 다음과 같습니다:

from dataclasses import dataclass

@dataclass
class Order:
    id: int
    customer: str
    total: float
    status: str

# 주문 목록
orders = [
    Order(1, "An", 250_000, "PAID"),
    Order(2, "Bình", 90_000, "CANCELLED"),
    Order(3, "An", 1_200_000, "PAID"),
    Order(4, "Chi", 450_000, "PENDING"),
]

# 모든 주문의 총액
totals = []
for order in orders:
    totals.append(order.total)
print(totals)

List comprehension으로 다시 작성하면:

totals = [order.total for order in orders]
print(totals)

[250000, 90000, 1200000, 450000]

구조는 [표현식 for 요소 in iterable]입니다. 왼쪽에서 오른쪽으로 읽으면: “orders의 각 order에 대해 order.total을 가져온다”입니다. 빈 list를 선언할 필요도 없고 append를 호출할 필요도 없습니다. “새로운 list를 만들고 있다”는 의도가 [ 기호에서 바로 드러납니다.

2.2. if를 사용한 요소 필터링

이미 결제된 주문만 가져오려면 끝에 if 절을 추가합니다:

paid_ids = [order.id for order in orders if order.status == "PAID"]
print(paid_ids)

[1, 3]

끝의 if는 필터 역할을 합니다. 조건을 만족하지 않는 요소는 제외되므로, 결과 list가 원본 list보다 짧을 수 있습니다.

2.3. 앞에 있는 if-else 표현식

필터링하는 대신 조건에 따라 모든 요소를 변환하려면, 표현식 부분에 if-else를 배치합니다:

labels = ["lớn" if order.total >= 500_000 else "nhỏ" for order in orders]
print(labels)

['nhỏ', 'nhỏ', 'lớn', 'nhỏ']

여기가 많은 사람들이 헷갈리는 부분입니다. 두 가지 규칙만 기억하면 됩니다:

• for 뒤의 if는 필터링에 사용됩니다. 이 절에는 else가 없습니다.
• for 앞의 a if cond else b는 조건부 표현식입니다. 이 절은 반드시 else가 있어야 하며 요소의 개수를 유지합니다.

[x for x in data if cond else y]를 작성하면 SyntaxError가 발생합니다.

3. Dict comprehension

Dict comprehension은 {} 기호와 key: value 쌍을 사용합니다. 예를 들어, id로 주문을 빠르게 조회해야 합니다:

orders_by_id = {order.id: order for order in orders}
print(orders_by_id[3])

Order(id=3, customer='An', total=1200000, status='PAID')

이미 있는 dict를 변환할 수도 있습니다. 예를 들어, 가격표에 10% 세금을 적용합니다:

prices = {"cafe": 30_000, "tra_sua": 45_000, "banh_mi": 20_000}
prices_with_vat = {name: round(price * 1.1) for name, price in prices.items()}
print(prices_with_vat)

{'cafe': 33000, 'tra_sua': 49500, 'banh_mi': 22000}

prices.items()는 (key, value) 쌍을 반환하고, for 루프 안에서 바로 unpack합니다.

유용한 팁은 dict를 뒤집는 것입니다:

status_codes = {"PAID": 1, "PENDING": 2, "CANCELLED": 3}
code_to_status = {code: status for status, code in status_codes.items()}
print(code_to_status[2])

PENDING

주의: 두 요소가 같은 key를 생성하면, 나중 요소가 무음으로 이전 요소를 덮어씁니다. 예를 들어 {order.customer: order.total for order in orders}는 “An”의 마지막 주문(1,200,000)만 유지하고, 250,000 주문은 경고 없이 사라집니다. 그룹화가 필요하면 vòng for와 함께 defaultdict를 사용하세요 (섹션 7 참조).

4. Set comprehension

Set comprehension도 {}를 사용하지만 한 개의 표현식만 있고 :가 없습니다. 결과는 자동으로 중복을 제거합니다. 예를 들어, 어떤 고객이 주문했는지 알고 싶다면:

customers = {order.customer for order in orders}
print(sorted(customers))

['An', 'Bình', 'Chi']

각 고객은 “An”이 두 개의 주문을 가지고 있어도 한 번만 나타납니다. set 내의 순서는 보장되지 않으므로, 위 예제는 출력할 때 sorted를 사용합니다.

Set comprehension은 빠른 멤버십 확인이 필요할 때 매우 유용합니다:

banned_emails = {"spam@x.com", "Fake@Y.com"}
normalized = {email.strip().lower() for email in banned_emails}
print("fake@y.com" in normalized)

True

주의: 빈 {}는 dict이지 set이 아닙니다. 빈 set을 만들려면 set()을 작성해야 합니다.

5. 여러 for 루프와 중첩된 comprehension

Comprehension 안에 여러 for를 작성할 수 있습니다. 읽는 순서는 위에서 아래로 중첩된 for 루프를 작성할 때와 정확히 같습니다:

carts = {
    "An": ["cafe", "banh_mi"],
    "Chi": ["tra_sua"],
}

items = [(customer, item) for customer, products in carts.items() for item in products]
print(items)

[('An', 'cafe'), ('An', 'banh_mi'), ('Chi', 'tra_sua')]

위의 comprehension은 다음 코드와 동치입니다:

items = []
for customer, products in carts.items():
    for item in products:
        items.append((customer, item))

가장 일반적인 경우는 중첩된 list를 평탄화(flatten)하는 것입니다:

matrix = [[1, 2, 3], [4, 5, 6]]
flat = [value for row in matrix for value in row]
print(flat)

[1, 2, 3, 4, 5, 6]

기억하는 방법: 가장 바깥쪽 for 루프가 먼저 나옵니다. 순서를 이해하기 위해 두 번째로 읽어야 한다면, 일반 for 루프로 나누는 것이 좋다는 신호입니다.

6. Generator expression: list가 필요 없을 때

[]를 ()로 바꾸면 generator expression을 얻습니다. Python은 메모리에 list를 생성하지 않습니다. 값은 코드가 소비할 때 하나씩 생성됩니다.

revenue = sum(order.total for order in orders if order.status == "PAID")
print(revenue)

1450000

Generator가 함수의 유일한 인수일 때, 괄호 한 쌍을 생략할 수 있습니다. 이 방식은 sum, any, all, max, min과 잘 맞습니다:

has_big_order = any(order.total > 1_000_000 for order in orders)
all_positive = all(order.total > 0 for order in orders)
print(has_big_order, all_positive)

True True

any는 조건을 만족하는 첫 번째 요소를 만나면 즉시 중지합니다. 따라서 대용량 데이터의 경우, generator는 메모리를 절약하면서도 any([...])보다 빠릅니다.

주의: Generator는 한 번만 순회할 수 있습니다. 두 번째 순회는 아무것도 반환하지 않습니다:

paid = (o for o in orders if o.status == "PAID")
print(len(list(paid)))
print(len(list(paid)))

2
0

결과를 여러 번 사용해야 한다면 list comprehension을 사용하세요.

7. Comprehension을 읽기 쉽게 유지하기

7.1. 복잡한 로직에 대해 함수의 이름 지정

긴 조건은 comprehension을 읽기 어렵게 만듭니다:

vip = [o.customer for o in orders if o.status == "PAID" and o.total >= 1_000_000 and not o.customer.startswith("test")]

조건을 명확한 이름을 가진 함수로 분리할 수 있습니다:

def is_vip_order(order: Order) -> bool:
    return (
        order.status == "PAID"
        and order.total >= 1_000_000
        and not order.customer.startswith("test")
    )

vip = [order.customer for order in orders if is_vip_order(order)]
print(vip)

['An']

이제 comprehension은 문장처럼 읽히고, 비즈니스 로직은 테스트할 별도의 위치가 있습니다.

7.2. 합리적인 줄바꿈

Comprehension은 한 줄에 있을 필요가 없습니다. 괄호 안에서 자유롭게 줄을 나눌 수 있습니다:

summary = {
    order.id: order.total
    for order in orders
    if order.status != "CANCELLED"
}
print(summary)

{1: 250000, 3: 1200000, 4: 450000}

각 줄은 한 역할을 합니다: 표현식, 소스, 필터.

7.3. 두 번 계산을 피하기 위한 Walrus 연산자

때로는 필터링과 동일한 계산의 결과를 모두 사용해야 합니다. := 연산자(Python 3.8+)는 한 번만 계산하도록 도와줍니다:

def discount_for(order: Order) -> float:
    return order.total * 0.05 if order.total >= 400_000 else 0

discounts = {
    order.id: d
    for order in orders
    if (d := discount_for(order)) > 0
}
print(discounts)

{3: 60000.0, 4: 22500.0}

walrus를 절제 있게 사용하세요. 변수 d는 comprehension 외부 범위로 “누출”됩니다. 또한 많은 독자가 이 구문에 익숙하지 않을 것입니다.

8. 데이터 변환 로직 테스트하기

Comprehension은 종종 데이터 변환 함수에 있으므로 다른 로직처럼 테스트해야 합니다. 먼저 pytest를 설치하세요:

pip install pytest

reports.py 모듈이 있다고 가정해봅시다:

# reports.py
from dataclasses import dataclass

@dataclass
class Order:
    id: int
    customer: str
    total: float
    status: str

def revenue_by_customer(orders: list[Order]) -> dict[str, float]:
    customers = {o.customer for o in orders if o.status == "PAID"}
    return {
        c: sum(o.total for o in orders if o.customer == c and o.status == "PAID")
        for c in customers
    }

그리고 해당 테스트 파일:

# test_reports.py
from reports import Order, revenue_by_customer

def test_revenue_only_counts_paid_orders():
    orders = [
        Order(1, "An", 250_000, "PAID"),
        Order(2, "An", 1_200_000, "PAID"),
        Order(3, "Bình", 90_000, "CANCELLED"),
        Order(4, "Chi", 450_000, "PAID"),
    ]

    result = revenue_by_customer(orders)

    assert result == {"An": 1_450_000, "Chi": 450_000}
    assert "Bình" not in result

def test_empty_orders_returns_empty_dict():
    assert revenue_by_customer([]) == {}

pytest -q

..                                                                       [100%]
2 passed in 0.01s

첫 번째 테스트는 두 가지를 주장합니다: 취소된 주문은 계산되지 않고, 취소된 주문만 있는 고객은 결과에 나타나지 않습니다. 두 번째 테스트는 빈 list라는 경계 케이스를 확인합니다.

위의 함수는 맞지만 O(n × 고객 수)의 복잡도를 가집니다. 각 고객이 전체 orders를 다시 순회하기 때문입니다. 대용량 데이터의 경우, for 루프와 defaultdict를 함께 사용하는 것이 더 빠르고 읽기 쉽습니다:

from collections import defaultdict

def revenue_by_customer(orders: list[Order]) -> dict[str, float]:
    totals: dict[str, float] = defaultdict(float)
    for order in orders:
        if order.status == "PAID":
            totals[order.customer] += order.total
    return dict(totals)

기존 두 테스트는 수정할 필요 없이 여전히 통과합니다. 이것이 테스트가 있는 것의 장점입니다: comprehension에서 for 루프로, 또는 그 반대로 자신 있게 리팩토링할 수 있습니다.

9. Comprehension을 사용하면 안 될 때

• Side effect가 필요할 때. [print(o) for o in orders] 또는 [db.save(o) for o in orders]를 작성하지 마세요. 이렇게 하면 의미 없는 None list가 생성되고 독자를 혼동시킵니다. for 루프를 사용하세요.
• 두 개 이상의 수준으로 중첩되거나 여러 if가 섞여 있을 때. 한 줄에 세 개의 for 루프는 거의 항상 줄을 나눈 버전보다 읽기 어렵습니다.
• 누계 또는 그룹화가 필요할 때, 예를 들어 합계 또는 개수 세기. defaultdict, Counter 또는 for 루프를 사용하세요.
• 각 요소에 대해 예외 처리가 필요할 때. Comprehension은 try/except를 포함할 수 없습니다. 별도의 함수로 분리하여 comprehension에서 호출하거나 for 루프를 사용해야 합니다.
• 한 번만 순회하는 매우 큰 데이터일 때. 이 경우 list 대신 generator expression을 사용하세요.

주의: Python 3에서는 comprehension의 루프 변수가 외부로 누출되지 않습니다. [x for x in range(3)] 이후, 이름 x는 외부 범위에 존재하지 않으며, 이는 일반 for 루프와 다릅니다. 유일한 예외는 섹션 7.3에서 언급했듯이 :=로 할당된 변수입니다.

10. 결론

Comprehension은 “기존 collection에서 새 collection을 만든다”는 의도를 간결하게 표현하는 데 도움이 됩니다:

• list는 [...], dict는 {k: v ...}, set은 {...}, generator는 (...).
• for 뒤의 if는 필터링에 사용되고, for 앞의 a if cond else b는 변환에 사용됩니다.
• 복잡한 조건을 명확한 이름의 함수로 분리하고, 표현식이 길 때는 줄을 나누세요.
• 메모리를 절약하기 위해 sum, any, all과 함께 generator expression을 사용하세요.
• Side effect가 있을 때, 그룹화가 필요할 때, 예외 처리할 때, 또는 comprehension이 읽기 어려워질 때 for 루프로 돌아가세요.

가장 간단한 규칙: 동료가 당신의 comprehension을 두 번 읽어야 이해한다면 for 루프로 다시 작성하세요. 간결한 코드는 읽기 쉬울 때만 가치가 있습니다.