Почему стоит использовать моржовый оператор в Python

Walrus-оператор (или моржовый оператор) существует в Python уже достаточно давно (начиная с версии 3.8), но до сих пор вызывает определённые споры и даже хейт.

В этой статье постараемся убедиться в том, что моржовый оператор действительно является хорошим дополнением к языку, и что если использовать его правильно, он может помочь сделать код более лаконичным и читабельным.

Основы

Если вы ещё не знакомы с :=, давайте рассмотрим некоторые из основных вариантов его использования.

Первый пример — это использование моржового оператора для уменьшения количества вызовов функций. Представим функцию под названием func(), которая выполняет очень дорогостоящие вычисления. Вычисление результатов занимает много времени, поэтому мы не хотим вызывать его много раз:

# "func" вызывана 3 раза
result = [func(x), func(x)**2, func(x)**3]

# Повторное использование результата "func" 
# без разбиения кода на несколько строк
result = [y := func(x), y**2, y**3]

В первом объявлении списка выше func(x) вызывается 3 раза, каждый раз возвращая один и тот же результат, что приводит к пустой трате времени и вычислительных ресурсов. При перезаписи с использованием моржового оператора функция func() вызывается только один раз, присваивая свой результат y и повторно используя его для вычисления оставшихся 2 значений списка. Вы могли бы сказать: "Я могу просто добавить y = func(x) перед объявлением списка, и мне не нужен моржовый оператор!", но это одна дополнительная, ненужная строка кода. Более того, если не знать, что функция func(x) тяжёлая, то существование строки с определением дополнительной переменной и вовсе может быть непонятным.

Если вас не убедило вышесказанное, то вот ещё один аргумент. Рассмотрим следующие варианты list comprehension с той же дорогостоящей func():

result = [func(x) for x in data if func(x)]

result = [y for x in data if (y := func(x))]

В первой строке func(x) вызывается дважды на каждой итерации. Во второй же мы вычисляем её один раз в операторе if, а затем повторно используем. Длина кода одинакова, обе строки одинаково читаемы, но вторая в два раза эффективнее. Можно избежать использования моржового оператора, сохраняя производительность, заменив его полноценным циклом for, но это потребует написания 5 строк кода.

Одним из наиболее распространенных вариантов использования моржового оператора является сокращение вложенных условий, например, при использовании регулярных выражений:

import re

test = "Something to match"

pattern1 = r"^.*(thing).*"
pattern2 = r"^.*(not present).*"

m = re.match(pattern1, test)
if m:
    print(f"Совпадение по первому шаблону: {m.group(1)}")
else:
    m = re.match(pattern2, test)
    if m:
        print(f"Совпадение по второму шаблону: {m.group(1)}")

# ---------------------
# а так чище
if m := (re.match(pattern1, test)):
    print(f"Совпадение по первому шаблону: '{m.group(1)}'")
elif m := (re.match(pattern2, test)):
    print(f"Совпадение по второму шаблону: '{m.group(1)}'")

Используя моржовый оператор, мы сократили код с 7 до 4 строк, сделав его более читаемым за счет удаления вложенных if.

Еще один пример упрощения кода — это "полуторный" идиоматический цикл:

while True:  # Цикл
    command = input("> ")
    if command == "exit":  # И половина
        break
    print("Была выбрана команда:", command)

# ---------------------
# а так чище
while (command := input("> ")) != "exit":
    print("Была выбрана команда:", command)

Обычным решением является использовать бесконечный цикл while и контроль выполнения делегировать оператору break. Вместо этого мы можем использовать моржовый оператор для переназначения значения command, а затем использовать его в условном цикле while в одной строке, что делает код намного чище и короче.

Аналогичное упрощение может быть применено и к другим циклам while, например, при построчном чтении файлов или при получении данных из сокета.

Расчеты нарастающим итогом

Переходим к некоторым немного более продвинутым вариантам использования моржового оператора. Это возможность делать расчёты нарастающим итогом:

data = [5, 4, 3, 2]
c = 0 
print([(c := c + x) for x in data])  # c = 14
# [5, 9, 12, 14]

from itertools import accumulate
print(list(accumulate(data)))

# ---------------------
data = [5, 4, 3, 2]
print(list(accumulate(data, lambda a, b: a*b)))
# [5, 20, 60, 120]

a = 1 
print([(a := a*b) for b in data])
# [5, 20, 60, 120]

Первые 3 строки показывают, как вы можете использовать моржовый оператор для вычисления текущего итога. Для такого простого случая лучше подходят функции из itertools, такие как accumulate, как мы видим в следующих 2 строках. Однако для более сложных сценариев использование itertools довольно быстро становится нечитаемым, а версия с := намного приятнее, чем версия с лямбдой.

Если всё ещё не убедились, посмотрите на примеры accumulate в документации (например, накопление процентов), которые довольно трудно читать. Попробуйте переписать их, используя моржовый оператор, и они будут выглядеть намного лучше.

Объявление переменных внутри f-строки

Этот пример скорее демонстрирует возможности и ограничения :=, а не лучшие практики. Но если уж очень хочется, то можно использовать его внутри f-строк:

from datetime import datetime

print(f"Today is: {(today:=datetime.today()):%Y-%m-%d}, which is {today:%A}")
# Today is: 2022-07-01, which is Friday

from math import radians, sin, cos

angle = 60
print(f'{angle=}\N{degree sign} {(theta := radians(angle)) =: .2f}, {sin(theta) =: .2f}, {cos(theta) =: .2f}')
# angle=60° (theta := radians(angle)) = 1.05, sin(theta) = 0.87, cos(theta) = 0.50

В первом примере выше := используется для определения переменной today, которая затем повторно используется в этой же строке, избегая повторного вызова datetime.today().

Аналогично, во втором примере объявляется переменная theta, которая затем повторно используется для вычисления sin(theta) и cos(theta). Тут ещё используется что-то, похожее на "обратного моржа" — это просто =, которое выводит выражение на одной строке со своим значением, а : используется для форматирования.

Обратите внимание, что все действия с моржовым оператором должны быть заключены в круглые скобки, чтобы f-строки могли правильно их интерпретировать.

Any и All

Вы можете использовать функции any() и all(), чтобы проверить, удовлетворяют ли какие-либо или все значения в итерируемом объекте определённому условию. Но что, если мы хотим также сохранить значение, вызвавшее возврат True функции any() (так называемый "свидетель") или значение, вызвавшее False в функции all() (так называемый "контрпример")?

numbers = [1, 4, 6, 2, 12, 4, 15]

# возвращаются логические значения
print(any(number > 10 for number in numbers))  # True
print(all(number < 10 for number in numbers))  # False

# ---------------------
any((value := number) > 10 for number in numbers)  # True
print(value)  # 12

all((counter_example := number) < 10 for number in numbers)  # False
print(counter_example)  # 12

Используя этот трюк, как any(), так и all() делают ленивую проверку. Это означает, что они останавливают вычисление, как только находят первый "свидетель" или "контрпример" соответственно. Поэтому переменная, созданная моржовым оператором, всегда даст нам первый "свидетель" / "контрпример".

Подводные камни и ограничения

Предыдущие примеры показывают достоинства моржового оператора, но важно перечислить его недостатки и ограничения. Вот несколько подводных камней, на которые можно наткнуться при его использовании.

В предыдущем примере мы видели, как использовать any()/all(), но в некоторых случаях это может привести к неожиданным результатам:

for i in range(1, 100):
    if (two := i % 2 == 0) and (three := i % 3 == 0):
        print(f"{i} кратно 6.")
    elif two:
        print(f"{i} кратно 2.")
    elif three:
        print(f"{i} кратно 3.")

# NameError: name 'three' is not defined

В примере выше создаётся условие с 2 присваиваниями, объединенными с помощью and, которые проверяют, делится ли число на 2, 3 или 6 на основе того, удовлетворяются ли первое, второе или оба условия.

На первый взгляд это может показаться хорошим трюком, но из-за ленивой проверки если выражение (two := i % 2 == 0) вернёт False, то вторая проверка вообще не будет делаться, и, следовательно, переменная three не будет определена или будут иметь неактуальное значение с предыдущей итерации цикла.

Такой подход можно применять осознано и с пользой, например, с регулярными выражениями для поиска нескольких шаблонов в строке:

import re

tests = ["Something to match", "Second one is present"]

pattern1 = r"^.*(thing).*"
pattern2 = r"^.*(present).*"

for test in tests:
    m = re.match(pattern1, test)
    if m:
        print(f"Совпадение по первому шаблону: {m.group(1)}")
    else:
        m = re.match(pattern2, test)
        if m:
            print(f"Совпадение по второму шаблону: {m.group(1)}")

# Совпадение по первому шаблону:: thing
# Совпадение по второму шаблону:: present

for test in tests:
    if m := (re.match(pattern1, test) or re.match(pattern2, test)):
        print(f"Совпадение: '{m.group(1)}'")
        # Совпадение: 'thing'
        # Совпадение: 'present'

Чуть выше уже был аналогичный пример, где использовалось if/elif в сочетании с моржовым оператором. Здесь же всё ещё проще, условие сводится if.

Если вы только знакомитесь с моржовым оператором, вы можете заметить, что он приводит к изменению поведения областей видимости переменных в comprehensions:

values = [3, 5, 2, 6, 12, 7, 15]

tmp = "unmodified"
dummy = [tmp for tmp in values]
print(tmp)  # Как и ожидалось, "tmp" не был перезаписан - он по-прежнему привязан к "unmodified"

total = 0
partial_sums = [total := total + v for v in values]
print(total)  # Вывод: 50

При обычном list/dict/set comprehension переменная в них не попадает во внешнюю область видимости, и поэтому любые существующие переменные с таким же именем не будут изменены. Однако, при использовании моржового оператора переменная из генератора (total) остается доступной после возврата генератора, принимая значение изнутри генератора.

Где никогда не стоит использовать моржовый оператор — это вместе с with:

class ContextManager:
    def __enter__(self):
        print("Входим в контекст..")

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Выходим из контекста..")

with ContextManager() as context:
    print(context)  # None

with (context := ContextManager()):
    print(context)  # <__main__.ContextManager object at 0x7fb551cdb9d0>

При использовании обычного синтаксиса with ContextManager() as context: ... context привязан к возвращаемому значению context.__enter__(), а если использовать вариант с :=, то он привязана к результату самого ContextManager(). Это часто не имеет значения, потому что context.__enter__() обычно возвращает self, но в случае, если это не так, это создаёт очень трудные для отладки проблемы.

Для более практического примера рассмотрим, что происходит, если использовать моржовый оператор с контекстным менеджером закрытия:

from contextlib import closing
from urllib.request import urlopen

with closing(urlopen('https://www.python.org')) as page:
    for line in page:
        print(line)  # Печатает содержимое HTML-файла сайта

with (page := closing(urlopen('https://www.python.org'))):
    for line in page:
        print(line)  # TypeError: 'closing' object is not iterable

Ещё одна проблема, с которой вы можете столкнуться — это приоритет :=, который ниже, чем у логических операторов:

text = "Something to match."
flag = True

if match := re.match(r"^.*(thing).*", text) and flag:
    print(match.groups())  # AttributeError: 'bool' object has no attribute 'group'

if (match := re.match(r"^.*(thing).*", text)) and flag:
    print(match.groups())  # ('thing',)

Здесь мы видим, что нам нужно заключить присваивание в круглые скобки, чтобы убедиться, что результат re.match(...) присваивается переменной. Если мы этого не сделаем, сначала вычисляется выражение and, и вместо него будет присвоен логический результат.

И ещё одно небольшое ограничение. Сейчас нельзя использовать однострочные type hints вместе с моржовым оператором. Поэтому, если вы хотите указать тип переменной, то нужно разделить ее на 2 строки:

from typing import Optional

value: Optional[int] = None
while value := some_func():
    ...  # Делаем что-либо

Заключение

Как и любой другой синтаксической возможностью, моржовым оператором можно злоупотреблять, что приведёт к снижению читаемости и лаконичности кода. Не нужно его пихать, куда угодно. Рассматривайте его как инструмент — понимайте его преимущества и недостатки и используйте его там, где это уместно.

Если вы хотите посмотреть ещё примеры использования данного оператора, то изучите, как он был введен в стандартную библиотеку CPython — все эти изменения можно найти в этом PR. Кроме того, можно ознакомиться с PEP 572, в котором есть ещё больше примеров, а также обоснование введения оператора.

Источник: Martin Heinz