Effective Python 컴프리헨션과 제너레이터 Better way 29

컴프리헨션(리스트, 딕셔너리, 집합 중 무엇이든)에서 같은 계산을 여러 위치에서 공유하는 경우가 흔하다. 예를 들어 한회사에서 주문을 관리하기 위한 프로그램을 만든다고 해볼게요.
고객이 새로운 주문을 보내면 주문을 처리할 만한 재고가 있는지 알려줘야 한다. 그러려면 고객의 요청이 재고 수량을 넘지 않고, 배송에 필요한 최소 수량(부품 수가 여덟 개 이상이어야 함)을 만족하는지 확인해야한다.

stock = {
	"못": 125,
    	"나사못": 35,
    	"나비너트": 8,
    	"와셔": 24
}

order = ["나사못", "나비너트", "클립"]

def get_batches(count, size):
	return count // size
    
result = []

for name in order:
	count = stock.get(name, 0)
    	batches = get_batches(count, 8)
        if batches:
        	result[name] = batches
            
print(result)
>>>
{"나사못": 4, "나비너트": 1}

여기서 딕셔너리 컴프리헨션을 사용하면 이 루프의 로직을 더 간결하게 표현 할 수 있다.

found = {name: get_batches(stock.get(name, 0), 8)

이 코드는 앞의 코드보다 짧지만 get_baches(stock.get(name, 0), 8)이 반복된다는 단점이 있다. 이로 인해 기술적으로는 불필요한 시각적인 잡음이 들어가서 가독성이 나빠진다. 그리고 두 식을 항상 똑같이 변경해야 하므로 실수할 가능성도 높아진다. 예를 들어 첫 번째 get_baches 호출에서만 두 번째 인자를 8 대신 4로 바꿨는데 결과가 달라진다.

has_bug = {name: get_baches(stock.get(name, 0), 4) 
	for name in order if get_batches(stock.get(name, 0), 8)}
    
print("예상:", found)
print("실제:", has_bug)

>>>
예상: {'나사못': 4, '나비너트': 1}
실제: {'나사못': 8, '나비너트': 2}

이러한 문제에 대한 쉬운 해법은 파이썬 3.8에 도입된 왈러스 연산자(:=)를 사용하는 것이다. 왈러스 연산자를 사용하면 컴프리헨션의 일부분에 대입식을 만들 수 있다.

found = {name: batches for name in order if (batches := get_batches(stock.get(name, 0), 8))}

대입식(batches := get_batches(...))을 사용하면 stock 딕셔너리에서 각 order 키를 한 번만 조회하고 get_batches를 한 번만 호출해서 그 결과를 batches 변수에 저장할 수 있다. 컴프리헨션의 다른 곳에서는 batches 변수를 참조해서 get_batches를 다시 호출할 필요 없이 딕셔너리의 내용을 만들 수 있다. get_batches를 얻기 위한 불필요한 함수 호출을 제거하면 order 리스트 안에 있는 각 원소에 대해 불필요한 연산을 수행하지 않으므로 성능도 향상된다.

대입식을 컴프리헨션의 값 식에 사용해도 문법적으로 올바르다. 하지만 컴프리헨션의 다른 부분에서 이 변수를 읽으려고 하면 컴프리헨션이 평가되는 순서 때문에 실행 시점에 오류가 발생할 것이다.

result = [name: (tenth := count //10) for name, count in stock.items() if tenth > 0 ]

>>>
Traceback ...
NameError: name 'tenth' is not defined

대입식을 조건 쪽으로 옮기고 대입식에서 만들어진 변수 이름을 컴프리헨션 값 식에서 참조하면 이 문제를 해결 할 수 있다.

result = [name: tenth for name, count in stock.items() if (tenth := count // 10 )> 0 ]

컴프리헨션이 값 부분에서 왈러스 연산자를 사용할 때 그 값에 대한 조건부분이 없다면 루프 밖 영역으로 루프 변수가 누출된다.

half = [(last := count //2) for count in stock.values()]
print(f'{hlf}의 마지막 원소는 {last}')

>>>
[62, 17, 4, 12]의 마지막 원소는 12

이러한 루프 변수 누출은 일반적인 for 루프에서 발생하는 루프 변수 누출과 비슷하다.

for count in stock.values(): # 루프 변수가 누출됨
	pass
print(f'list(stock.values())}의 마지막 원소는 {count}')

>>>
[125, 35, 8, 24]의 마지막 원소는 24

하지만 컴프리헨션의 루프 변수인 경우에는 비슷한 누출이 생기지 않는다. 

# 바로 앞의 예제를 처리하다가 count가 정의된 경에에는 제대로 작동하지 않음
# 파이썬을 재시작하고 아래 코드를 실행해야 오류를 볼 수 있음
half = [count // 2 for count in stock.values()]
print(half) # OK
print(count) # 루프 변수가 누출되지 않기 때문에 예외가 발생

>>>
[]62, 17, 4, 12]
Traceback ...
NameError: name 'count' is not defined

루프 변수를 누출하지 않는 편이 낫다. 따라서 컴프리헨션에서 대입식을 조건에만 사용하는 것을 권장한다.

대입식은 제너리에터의 경우에도 똑같은 방식으로 동작한다.

아래 코드에서는 딕셔너리 인스턴스 대신 제품 이름과 현재 재고 수량 쌍으로 이뤄진 이터레이터를 만든다.

found = ((name, batches) for name in order 
			 if (batches := get_batches(stock.get(name, 0), 8)))
 print(next(found))
 print(next(found))
 
 >>> 
 ('나사못', 4)
 ('나비너트', 1)

요약

• 대입식을 통해 컴프리헨션이나 제너레이터의 식의 조건 부분에서 사용한 값을 같은 컴프리헨션이나 제너레이터의 다른 위치에서 재사용할 수 있다. 이를 통해 가독성과 성능을 향상시킬 수 있다.
• 조건이 아닌 부분에도 대입식을 사용할 수 있지만, 그런 사용은 피해야 한다.