파이썬 컴프리헨션으로 데이터를 우아하게 다루기

파이썬 코드를 읽다 보면 대괄호 안에 for가 들어간 독특한 문법을 자주 만나게 됩니다. 처음 보면 낯설지만, 한 번 손에 익으면 도저히 안 쓸 수가 없는 파이썬의 꽃 🌼이라고 할 수 있는데요. 바로 컴프리헨션(comprehension)입니다.

컴프리헨션은 기존 데이터를 기반으로 새로운 리스트, 딕셔너리, 세트를 간결하게 만들어내는 파이썬만의 문법입니다. 반복문을 여러 줄에 걸쳐 쓰는 대신 한 줄로 표현할 수 있어서 코드가 훨씬 깔끔해지거든요.

이번 포스팅에서는 리스트 컴프리헨션부터 시작해서 딕셔너리, 세트 컴프리헨션, 그리고 제너레이터 표현식까지 하나씩 다뤄보겠습니다.

리스트 컴프리헨션 기본

리스트 컴프리헨션(list comprehension)은 기존 리스트를 변환하여 새 리스트를 만드는 문법입니다. 기본 형태는 [표현식 for 변수 in 반복가능객체]입니다.

예를 들어 1부터 5까지의 제곱을 구한다고 해볼게요.

먼저 일반적인 for 루프로 작성하면 이렇습니다.

>>> squares = []
>>> for n in range(1, 6):
...     squares.append(n ** 2)
...
>>> squares
[1, 4, 9, 16, 25]

빈 리스트를 만들고, 반복문을 돌면서 append()로 하나씩 추가하는 패턴이죠. 동작에는 문제가 없지만 뭔가 장황합니다.

리스트 컴프리헨션을 쓰면 이걸 한 줄로 줄일 수 있습니다.

>>> [n ** 2 for n in range(1, 6)]
[1, 4, 9, 16, 25]

for 루프 앞에 있는 n ** 2가 각 원소에 적용할 표현식이고, for n in range(1, 6)이 반복 부분입니다. 대괄호로 감싸면 결과가 바로 리스트로 나오기 때문에 별도로 list()를 호출하지 않아도 됩니다.

문자열을 다룰 때도 유용합니다.

>>> names = ["alice", "bob", "charlie"]
>>> [name.upper() for name in names]
['ALICE', 'BOB', 'CHARLIE']

>>> words = ["Hello", "World", "Python"]
>>> [len(word) for word in words]
[5, 5, 6]

각 원소에 메서드를 호출하거나 함수를 적용하는 건 물론이고, 어떤 표현식이든 넣을 수 있어서 활용 범위가 넓습니다.

조건부 필터링

리스트 컴프리헨션에 if를 추가하면 특정 조건을 만족하는 원소만 골라낼 수 있습니다. 형태는 [표현식 for 변수 in 반복가능객체 if 조건]입니다.

짝수만 골라내는 코드를 비교해보겠습니다.

>>> # for 루프 버전
>>> evens = []
>>> for n in range(10):
...     if n % 2 == 0:
...         evens.append(n)
...
>>> evens
[0, 2, 4, 6, 8]

>>> # 컴프리헨션 버전
>>> [n for n in range(10) if n % 2 == 0]
[0, 2, 4, 6, 8]

뒤쪽의 if n % 2 == 0 부분이 필터 역할을 합니다. 조건이 참인 원소만 최종 리스트에 포함되는 거죠.

실무에서 자주 쓰이는 패턴 몇 가지를 더 살펴볼까요?

>>> scores = [85, 42, 91, 67, 55, 78, 93, 38]
>>> [s for s in scores if s >= 70]
[85, 91, 67, 78, 93]

>>> files = ["report.pdf", "data.csv", "image.png", "notes.txt", "backup.csv"]
>>> [f for f in files if f.endswith(".csv")]
['data.csv', 'backup.csv']

>>> mixed = [1, "hello", 3.14, None, True, "world", 42]
>>> [x for x in mixed if isinstance(x, str)]
['hello', 'world']

이런 식으로 조건을 걸어서 원하는 데이터만 깔끔하게 뽑아낼 수 있습니다.

참고로, map()이나 filter() 같은 내장 함수를 사용해도 비슷한 결과를 얻을 수 있는데요. 컴프리헨션이 더 파이썬다운(pythonic) 방식으로 여겨지는 이유는 별도의 lambda가 필요 없고, 변환과 필터링을 하나의 표현식 안에서 처리할 수 있기 때문입니다.

조건부 표현식

필터링과 헷갈리기 쉬운데, if를 for 앞쪽에 두면 각 원소의 값을 조건에 따라 다르게 변환할 수 있습니다. 이때는 반드시 else와 함께 써야 합니다.

>>> numbers = [1, 2, 3, 4, 5]
>>> ["짝수" if n % 2 == 0 else "홀수" for n in numbers]
['홀수', '짝수', '홀수', '짝수', '홀수']

여기서 "짝수" if n % 2 == 0 else "홀수"는 파이썬의 삼항 표현식(ternary expression)입니다. for 뒤에 오는 if는 필터(포함할지 말지), for 앞에 오는 if-else는 변환(어떤 값으로 바꿀지)이라는 점을 기억하면 혼동을 피할 수 있습니다.

둘을 조합해서 쓸 수도 있는데요.

>>> nums = range(-5, 6)
>>> ["양수" if n > 0 else "0" for n in nums if n >= 0]
['0', '양수', '양수', '양수', '양수', '양수']

이 코드는 먼저 if n >= 0으로 음수를 걸러낸 다음, 나머지 원소에 대해 양수와 0을 구분합니다. 조건이 여러 개 겹치면 읽기 어려워질 수 있으니 적당한 선에서 사용하는 게 좋습니다.

중첩 반복

컴프리헨션 안에 for를 여러 개 쓰면 중첩 루프와 같은 효과를 낼 수 있습니다.

예를 들어, 구구단 일부를 만든다고 해봅시다.

>>> # 일반 for 루프
>>> result = []
>>> for x in range(2, 4):
...     for y in range(1, 4):
...         result.append(f"{x} x {y} = {x * y}")
...
>>> result
['2 x 1 = 2', '2 x 2 = 4', '2 x 3 = 6', '3 x 1 = 3', '3 x 2 = 6', '3 x 3 = 9']

>>> # 컴프리헨션 버전
>>> [f"{x} x {y} = {x * y}" for x in range(2, 4) for y in range(1, 4)]
['2 x 1 = 2', '2 x 2 = 4', '2 x 3 = 6', '3 x 1 = 3', '3 x 2 = 6', '3 x 3 = 9']

for를 나열한 순서가 바깥 루프에서 안쪽 루프 순서라는 점만 기억하면 됩니다. 왼쪽의 for x가 바깥 루프, 오른쪽의 for y가 안쪽 루프입니다.

2차원 리스트를 1차원으로 펼치는 작업도 중첩 컴프리헨션으로 깔끔하게 처리할 수 있습니다.

>>> matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
>>> [n for row in matrix for n in row]
[1, 2, 3, 4, 5, 6, 7, 8, 9]

다만 중첩이 두 단계를 넘어가면 가독성이 급격히 떨어집니다. 그럴 때는 일반 for 루프를 쓰는 편이 동료 개발자에게 훨씬 친절한 코드가 됩니다 😅

딕셔너리 컴프리헨션

자, 이제 리스트를 벗어나 봅시다. 중괄호({}) 안에 키: 값 형태로 쓰면 딕셔너리 컴프리헨션이 됩니다.

>>> {n: n ** 2 for n in range(1, 6)}
{1: 1, 2: 4, 3: 9, 4: 16, 5: 25}

두 개의 리스트를 합쳐서 딕셔너리를 만들 때 특히 유용합니다.

>>> keys = ["name", "age", "city"]
>>> values = ["Alice", 30, "Seoul"]
>>> {k: v for k, v in zip(keys, values)}
{'name': 'Alice', 'age': 30, 'city': 'Seoul'}

zip() 함수가 낯선 분은 파이썬 zip 함수 포스팅을 참고하세요.

기존 딕셔너리에서 조건에 맞는 항목만 추려내는 것도 가능합니다.

>>> prices = {"apple": 1200, "banana": 500, "cherry": 3000, "grape": 2500}
>>> {fruit: price for fruit, price in prices.items() if price >= 2000}
{'cherry': 3000, 'grape': 2500}

키와 값을 뒤집는 것도 한 줄이면 됩니다.

>>> original = {"a": 1, "b": 2, "c": 3}
>>> {v: k for k, v in original.items()}
{1: 'a', 2: 'b', 3: 'c'}

파이썬 딕셔너리를 다루는 실무 코드에서 컴프리헨션이 빠지는 경우가 거의 없을 정도로 자주 쓰이는 패턴이니 꼭 익혀두시길 바랍니다.

세트 컴프리헨션

중괄호를 쓰되 키: 값이 아니라 값만 넣으면 어떻게 될까요? 세트가 만들어집니다.

>>> {n % 3 for n in range(10)}
{0, 1, 2}

0부터 9까지의 숫자를 3으로 나눈 나머지를 세트로 모았더니 중복이 제거되어 {0, 1, 2}만 남았습니다.

문자열에서 사용된 모음만 뽑아내는 것도 간단합니다.

>>> sentence = "Hello World Python Programming"
>>> {ch.lower() for ch in sentence if ch.lower() in "aeiou"}
{'e', 'a', 'i', 'o'}

세트의 특성상 중복이 자동으로 제거되기 때문에 고유한 값을 모을 때 아주 편리합니다.

제너레이터 표현식

마지막으로 소개할 형태가 있습니다. 대괄호 대신 소괄호(())로 감싸면 제너레이터 표현식(generator expression)이라는 전혀 다른 물건이 나옵니다.

>>> gen = (n ** 2 for n in range(1, 6))
>>> gen
<generator object <genexpr> at 0x104a8f6d0>

문법은 리스트 컴프리헨션과 거의 같지만 동작 방식이 완전히 다릅니다. 리스트 컴프리헨션은 모든 원소를 한꺼번에 메모리에 올리는 반면, 제너레이터 표현식은 원소를 하나씩 필요할 때마다 생산합니다.

>>> gen = (n ** 2 for n in range(1, 6))
>>> next(gen)
1
>>> next(gen)
4
>>> next(gen)
9

next()를 호출할 때마다 다음 값을 계산해서 돌려주는 거죠. 전체 결과를 리스트로 변환하고 싶다면 list()로 감싸면 됩니다.

>>> list(n ** 2 for n in range(1, 6))
[1, 4, 9, 16, 25]

제너레이터 표현식이 진짜 빛을 발하는 순간은 대용량 데이터를 처리할 때입니다.

>>> # 리스트 컴프리헨션: 천만 개를 한꺼번에 메모리에 올림
>>> import sys
>>> sys.getsizeof([n for n in range(10_000_000)])
89095160

>>> # 제너레이터 표현식: 메모리를 거의 사용하지 않음
>>> sys.getsizeof(n for n in range(10_000_000))
200

약 89MB와 200바이트, 거의 45만 배 차이입니다! 데이터가 커질수록 이 차이는 더 벌어집니다.

sum(), max(), min(), any(), all() 같은 내장 함수에 제너레이터 표현식을 바로 넘기는 패턴도 자주 사용됩니다.

>>> sum(n ** 2 for n in range(1, 11))
385
>>> max(len(word) for word in ["apple", "banana", "cherry"])
6
>>> any(n > 100 for n in [3, 50, 7, 120, 9])
True

이때 함수 호출의 소괄호가 이미 있으므로 제너레이터 표현식을 위한 소괄호를 따로 쓰지 않아도 됩니다.

왈러스 연산자 활용

파이썬 3.8에서 도입된 왈러스 연산자(:=)를 컴프리헨션과 함께 쓰면 계산 결과를 변수에 담으면서 동시에 조건 검사를 할 수 있습니다.

>>> data = ["hello", "", "world", "", "python"]
>>> [upper for s in data if (upper := s.upper())]
['HELLO', 'WORLD', 'PYTHON']

s.upper()를 upper에 할당하면서 빈 문자열(falsy)은 걸러내고, 결과값을 바로 사용합니다. 왈러스 연산자가 없었다면 s.upper()를 두 번 호출하거나 별도 변수를 써야 했겠죠.

좀 더 실용적인 예를 보겠습니다.

>>> import re
>>> texts = ["Phone: 010-1234-5678", "No number here", "Call 02-987-6543"]
>>> [m.group() for t in texts if (m := re.search(r"\d{2,3}-\d{3,4}-\d{4}", t))]
['010-1234-5678', '02-987-6543']

정규 표현식 매치를 시도하면서 매치된 경우에만 결과를 가져오는 패턴입니다. 이런 식으로 비용이 큰 연산의 결과를 재사용해야 할 때 왈러스 연산자가 빛을 발합니다.

컴프리헨션을 피해야 할 때

컴프리헨션은 강력하지만 만능은 아닙니다. 무조건 한 줄로 줄이겠다는 욕심이 오히려 코드를 읽기 어렵게 만들기도 합니다.

이럴 때는 일반 for 루프가 더 나은 선택입니다.

우선 로직이 복잡해서 한 줄이 길어지는 경우입니다.

>>> # 읽기 어려운 컴프리헨션 😵
>>> result = [transform(x) for x in data if validate(x) and x.status == "active" and x.score > threshold]

>>> # 일반 루프가 더 명확 👍
>>> result = []
>>> for x in data:
...     if not validate(x):
...         continue
...     if x.status != "active":
...         continue
...     if x.score <= threshold:
...         continue
...     result.append(transform(x))

또한 부수 효과(side effect)가 목적인 경우에도 컴프리헨션은 적합하지 않습니다.

>>> # 이렇게 하지 마세요 🚫
>>> [print(item) for item in items]

print()를 호출하면서 쓸모없는 None 리스트를 만들어내는 코드입니다. 이런 경우는 그냥 for 루프를 쓰는 게 맞습니다.

>>> # 이렇게 하세요 ✅
>>> for item in items:
...     print(item)

마지막으로, 중첩이 세 단계 이상이 되면 본인조차 나중에 코드를 이해하기 힘들어집니다. 경험상 for가 두 개를 넘어가면 일반 루프로 풀어 쓰는 게 장기적으로 유지보수하기 좋습니다.

마치며

이번 포스팅에서는 리스트 컴프리헨션의 기본 문법부터 시작해서 딕셔너리, 세트, 제너레이터 표현식까지 쭉 훑어봤습니다.

컴프리헨션의 핵심은 “간결함”입니다. 코드를 짧게 만드는 것이 목적이 아니라, 의도를 명확하게 드러내는 것이 진짜 목적입니다. 한 줄로 표현했을 때 오히려 의도가 흐려진다면 과감하게 일반 루프를 선택하세요.

map()이나 filter() 같은 함수형 도구와 비교해서 어느 쪽이 더 낫다고 단정 짓기는 어렵지만, 파이썬 커뮤니티에서는 대체로 컴프리헨션을 더 파이썬다운 방식으로 선호하는 편입니다. 두 가지 접근법을 모두 알아두고 상황에 맞게 골라 쓰는 게 가장 좋겠죠 😊

더 자세한 내용은 파이썬 공식 튜토리얼의 Data Structures 문서를 참고하세요.