29. 대입식을 사용해 컴프리헨션 안에서 반복 작업을 피하라¶

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

In [1]:

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

In [2]:

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

In [3]:

def get_batches(count, size):
    return count // size

In [5]:

result = {}
for name in order:
    count = stock.get(name, 0)
    batches = get_batches(count, 8)
    if batches:
        result[name] = batches

print(result)

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

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

In [8]:

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}

이 코드는 앞의 코드보다 짧지만 반복된다는 단점이 있다.

아래와 같은 실수를 할 수 있다.

In [9]:

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

{'나사못': 8, '나비너트': 2}

이러한 문제에 대한 쉬운 해법은 파이썬 3.8에 도입된 왈러스 연산자(:=)를 사용하는 것이다.

In [10]:

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

In [11]:

found

Out[11]:

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

In [12]:

batches

Out[12]:

대입식을 컴프리헨션의 값 식에 사용해도 문법적으로 올바르다.

하지만 컴프리헨션의 다른 부분에서 이 변수를 읽으려고 하면 컴프리헨션이 평가되는 순서 때문에 실행 시점에 오류가 발생할 것이다.

In [13]:

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

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-13-fc56e5c9d356> in <module>
----> 1 result = {name: (tenth := count // 10)
      2           for name, count in stock.items() if tenth > 0}

<ipython-input-13-fc56e5c9d356> in <dictcomp>(.0)
      1 result = {name: (tenth := count // 10)
----> 2           for name, count in stock.items() if tenth > 0}

NameError: name 'tenth' is not defined

대입식을 조건 쪽으로 옮기면 사용 가능하다

In [14]:

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

In [15]:

print(result)

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

In [16]:

tenth

Out[16]:

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

In [17]:

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

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

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

In [19]:

for count in stock.values():
    pass
print(f'{list(stock.values())}의 마지막 원소는 {count}')

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

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

In [20]:

half = [count // 2 for count in stock.values()]

In [21]:

print(half)

[62, 17, 4, 12]

In [23]:

print(count) # 원래 에러나옴

따라서 컴프리헨션에서 대입식을 조건에만 사용하는 것을 권장한다.

제너레이터의 경우에도 똑같다

In [24]:

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

In [25]:

print(next(found))

('나사못', 4)

In [26]:

print(next(found))

('나비너트', 1)

In [27]:

found

Out[27]:

<generator object <genexpr> at 0x7feeb760a4a0>

기억해야 할 내용¶

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