Pythonのセイウチ演算子 (:=) について調べた

セイウチ演算子 := は Rust における if let に相当する演算子で,最近 Python の調べものをしてたときにコード例でたまたま見かけて「この演算子 Python 3.8 からあったの!?」とひっくり返ってしまったのでメモ.

使い方

セイウチ演算子は ifwhile 文の条件式や for 文のイテレータ式の中で代入を実行できる構文.

if文の例

if文では以下のように使う:

a = [1, 5, -3, 0, 2, 2, 4, 3, -1, 0, 0]

if (n := len(a)) > 10:
    print(f'List is too long ({n} elements, expected <= 10)')

この例では,リスト a の長さを関数 len によって取得し,その値を if の条件式とブロックの中の2箇所で使う必要がある.

セイウチ演算子が無い場合は,

  • n = len(a)if の直前に書く
  • len(a) の呼び出しを2度書く

のどちらかを行う必要がある.前者はこの if 文以外で n を参照しない場合に冗長な書き方になるし,後者だと len(a) が何度も呼び出したくない重たい処理の場合に困るが,セイウチ演算子を使った書き方はコードの複雑さを減らしつつ可読性を向上させているのがわかる.

for文 + tqdmの例

tqdm/tqdm: :zap: A Fast, Extensible Progress Bar for Python and CLI

tqdm はfor文などで使うイテレータにプログレスバーを表示するモジュール.例えば以下ように書くことで,for文によって10回実行する処理にプログレスバーをつけることができる.

import time
from tqdm import tqdm


for k in tqdm(range(10)):
    time.sleep(0.1)

また,tqdm のインスタンスについて set_postfix などのメソッドを呼び出すことでプログレスバーの表示内容をカスタマイズすることができる. PyTorchなどを使った最適化の処理を書く際,目的関数の値をターミナル上でざっくり確認する際によく使う.

セイウチ演算子を使わずに set_postfix を利用しようとすると以下のような書き方になる:

import time
from tqdm import tqdm


pbar = tqdm(range(10))

for k in pbar:
    time.sleep(0.1)
    pbar.set_postfix({'k': k})

for文で使うイテレータの定義をfor文の外に書かないといけなくなるので少し読みにくく感じる.

セイウチ演算子を使うと以下のようになる:

import time
from tqdm import tqdm


for k in (pbar := tqdm(range(10))):
    time.sleep(0.1)
    pbar.set_postfix({'k': k})

コード量としては一行減っただけだが,やりたいことの本筋でない処理による複雑さが減るのはすごく気分が良い.

参考