[Python] ':=' 왈러스 연산자로 대입식을 사용한 컴프리헨션

컴프리헨션에서 같은 계산을 위해 여러 위치에서 공유하는 경우가 흔합니다.

아래의 예씨는 회사에서 주문을 관리하기 위한 프로그램이고 재고를 확인하는 과정입니다.

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}

# Dictionary Comprehension
found = {name: get_batches(stock.get(name, 0), 8)
         for name in order
         if get_batches(stock.get(name, 0), 8)}
print(found)

{'나사못': 4, '나비너트': 1}

앞의 코드 보다는 짧으나 get_batches가 반복된다는 단점이 있습니다. 두 식을 항상 똑같이 변경해야 합니다. 첫 번째 get_batches 에서 두 번째 인자를 8에서 4로 바꿀 경우, 결과가 달라집니다.

has_bug = {name: get_batches(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}

이러한 문제의 해결점은 := 왈러스 연산자를 사용하는 것입니다. 왈러스 연산자는 컴프리헨션의 일부분에 대입식을 만듭니다.

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

다시 두 코드를 모아놓고 비교하면, := 의 사용으로 훨씬 간편해졌음을 알 수 있습니다.

# 기존 코드
has_bug = {name: get_batches(stock.get(name, 0), 4)
           for name in order
           if get_batches(stock.get(name, 0), 8)}

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

**:= 의 사용으로 달라진 점**

1. order 키와 get_batches를 한 번만 조회하여 그 결과를 변수 batches 에 저장할 수 있습니다.
2. get_batches 를 다시 호출할 필요 없이 딕셔너리를 만들 수 있습니다.
3. 불필요한 함수 호출을 제거하여 불필요한 연산을 수행하지 않아 성능이 향상됩니다.

에러가 발생하는 코드

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

NameError: name 'tenth' is not defined

오류가 발생하는 부분은 어디일까요?

바로 tenth := count // 10 가 아닌 if tenth > 0 의 코드에서 발생합니다. if 절은 for … 과 변수 영역이 같습니다. 그러나 tenth는 for … 내부에 정의돼 있지 않으므로 값을 읽을 때 오류가 발생합니다.

result = {name: tenth for name, count in stock.items()
          if (tenth := count // 10) > 0}
print(result)

{'못': 12, '나사못': 3, '와셔': 2}

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

누출

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

half = [(last := count // 2) for count in stock.values()]
print(f'{half}의 마지막 원소는 {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)  # 작동함
print(count) # 루프 변수가 누출되지 않기 때문에 예외가 발생함

[62, 17, 4, 12]
24 # 제대로 작동시

Traceback … # 오류 발생시

NameError: name ‘count’ is not defined

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

대입식은 제너레이터의 경우에도 동일하게 작동합니다.

딕셔너리 인스턴스 대신 제품 이름과 현재 재고 수량 상으로 이뤄진 이터레이터를 만듭니다.

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

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

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

('나사못', 4)
('나비너트', 1)

요약

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