Effective Python C1. Pythonic(10) - 왈러스 연산자의 사용으로 대입문을 간소화

Introduction

왈러스 연산자(walrus operator)는 python 3.8부터 도입된 연산자로, 코드 중복 문제를 해결하기 위해 사용된다.

# 일반 대입문
a = b
# 왈러스 연산자
a := b

왈러스 연산자는 대입문이 쓰일 수 없는 위치에서 변수에 값을 대입해 조건문 등으로 사용할 수 있다. 대입식의 값은 왈러스 연산자 왼쪽에 있는 식별자에 대입된 값으로 평가된다.

Best WAY 10 : 대입식을 사용해 반복을 피하라

할당과 조건문을 한 줄로

예를 들어 아래와 같이 바구니에서 레몬이 몇 개인지 세고 그만큼 레모네이드를 만드는 코드를 짠다고 하자.

fresh_fruit = {
    'apple': 10,
    'banana': 5,
    'lemon': 2,
}
count = fresh_fruit.get('lemon', 0)
if count:
    make_lemonade(count)
else:
    out_of_stock()

count값이 fresh_fruit.get에 의해 할당된 뒤 바로 조건문에 쓰이기 때문에 이를 한 줄로 줄이면 좋을 것 같다. 그러나 기본 대입 연산자를 if문에 넣을 수는 없고, if문에 그냥 fresh_fruit.get(‘lemon’, 0)를 때려넣자니 if 문 안쪽에서 쓰인다.

아래와 같이 왈러스 연산자로 대입문과 if문에 의한 조건절을 하나로 합칠 수 있다.

if count := fresh_fruit.get('lemon', 0):
    make_lemonade(count)
else:
    out_of_stock()

또는 아래와 같이 비교문을 쓸 수도 있다.

if (count := fresh_fruit.get('lemon', 0)) >= 4:
    make_lemonade(count)
else:
    out_of_stock()

다중 선택 구문과 유사한 코드로 작성이 가능

C언어의 switch/case와 같은 다중 선택 구문이 필요할 경우, 좀 더 가독성있게 짤 수 있다.
예를 들어 바나나 > 레몬 > 사과 순으로 개수를 체크해 주스를 만든다고 하자.

기존 대입식을 쓸 경우, 아래와 같이 작성을 할 수도 있다.

count = fresh_fruit.get('banana', 0)
if count >= 2:
    pieces = slice_bananas(count)
else:
    count = fresh_fruit.get('lemon', 0)
    if count >= 4:
        pieces = slice_lemons(count)
    else: 
        count = fresh_fruit.get('apple', 0)
        if count >= 2:
            pieces = slice_apples(count)     
        else:
            pieces = 0
to_enjoy = make_juice(pieces)

왈러스 연산자를 사용하면 아래와 같이 쓸 수 있다.

if (count := fresh_fruit.get('banana', 0) >= 2):
    pieces = slice_bananas(count)
elif (count := fresh_fruit.get('lemon', 0) >= 4):
    pieces = slice_lemons(count)
elif (count := fresh_fruit.get('apple', 0) >= 2):
    pieces = slice_apples(count)
else:
    pieces = 0
to_enjoy = make_juice(pieces)

do-while 과 유사한 코드 작성이 가능.

가지고 있는 신선한 과일을 모두 주스로 만드는 로직을 구현해보자.

bottles = []
fresh_fruit = pick_fruit()

while fresh_fruit:
    for fruit, count in fresh_fruit.items():
        batch = make_juice(fruit, count)
        bottle.extend(batch)
    fresh_fruit = pick_fruit()

do-while문이 없어서 fresh_fruit 값을 일단 pick_fruit으로 할당받고, 그걸 주스로 만들어 병에 담고 있다.

왈러스 연산자는 대입과 조건문을 동시에 한줄로 만들 수 있어 위 상황을 마치 do-while문 처럼 구현이 가능하다.

bottles = []
while fresh_fruit := pick_fruit():
    for fruit, count in fresh_fruit.items():
        batch = make_juice(fruit, count)
        bottle.extend(batch)
        

Discussion

• 왈러스 연산자는 좋아보이긴 하는데 python 3.8부터 사용이 가능한 점이 아쉽다.

Summary

• 왈러스 연산자(:=)를 사용해 변수의 대입식과 변수 값의 참조(조건문 등으로)를 한번에 쓸 수 있다.

Reference

파이썬 코딩의 기술 제 2판. - 브렛 슬리킨 지음 / 오현석 옮김