Enum в стиле Rust-а в Python / Хабр

TL;DR

Реализация энумов в стиле раста в питоне: rust_enum.

Проблема

Сегодня я работал над небольшим ИИ модулем для моего личного геймдев-проекта и повстречал проблему с магическими константами. Вот, смотрите:

def try_producing_target(self, subject, perception):
    if self.subject is None: 
        return False  # what does False mean?
    
    if abs2(sub2(subject.position, self.subject.position)) <= self.distance: 
        return None  # what is the difference with None?

    if self.period.step():
        return self.subject.position
      
    return False  # False again, wtf?

Этот метод реализует ИИ компоненту, позволяющую NPC поменять свой путь и следовать за каким-то другим NPC или игроком. Если не установлено, что надо за кем-то следить, не надо менять текущую цель пути; если цель близка, нужно остановиться; если всё хорошо, надо менять цель пути на позицию NPC, за которым мы следуем, каждые N тиков (для оптимизации).

Единственная проблема с этим кодом -- это возвращаемые значения. Что означает False? А None? По сути это магические константы, не несущие никакого значения для потенциального читателя этого кода. С другой стороны ситуация ничуть не лучше:

def make_decision(self, subject, perception):
    # noinspection PySimplifyBooleanCheck
    if (target := self.follower.try_producing_target(subject, perception)) != False: 
        self.pather.going_to = target
      
    if (action := self.pather.try_going(subject, perception)) is not None: 
        return action

Это довольно простая функция в небольшом модуле, так что пока ни о чём можно не переживать; с другой стороны, в ходе разработки модуль вырастет, и вполне вероятно что я повстречаю эту проблему снова и снова и это уже ощутимо попортит общую читаемость. Чтобы решить эту проблему, в идеале нужно объявить ограниченный набор значений, который может возвращать функция, каждое с понятным именем. Идеальное решение в такой ситуации -- enum-ы в Rust:

enum TargetChange {
  Nothing,
  To(Option<int2>),
}

Когда мы устанавливаем тип возвращаемого значения try_producing_target на этот enum, это говорит компилятору и читателю, что функция возвращает либо что изменений цели не требуется, либо что надо поменять цель на определённое значение. Никаких магических значений, только константы с осмысленными именами, сгруппированные в один enum. Также, функция не может возвращать никаких других значений и полученный результат можно удобно засунуть в match. Получается очень читаемый и красивый код, и это было бы идеальным решением для такого класса проблем, так что наверняка в PyPI есть библиотека которая это всё делает, да ведь?

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

Решение

Enum в стиле Rust скорее всего должен использовать синтаксис классов и быть помеченным декоратором, показывающим что это enum, что-то вот такое:

@enum
class TargetChange:
    Nothing = {}
    To = {"target": Optional[int2]}

Самый простой способ подружить enum с match-ем это динамически создавать датаклассы для каждого атрибута, не являющегося дандерметодом:

from dataclasses import make_dataclass


def enum(cls):
    for field_name in dir(cls):
        if field_name.startswith('__') and field_name.endswith('__'): 
            continue
        
        setattr(cls, field_name, make_dataclass(
          field_name, list(getattr(cls, field_name).items()), bases=(cls, )
        ))
    return cls

Это работает, но к сожалению TargetChange.To(...) вызывает ворнинги, поскольку линтер думает что мы пытаемся вызвать словарь; также если мы создаём датаклассы из каждого атрибута, мы не можем добавлять пользовательские атрибуты и методы в enum. Так что новый синтаксис:

@enum
class TargetChange:
    Nothing = Case()
    To = Case(target=Optional[int2])

И лучшая реализация:

from dataclasses import make_dataclass


def enum(cls):
    for field_name in dir(cls):
        if not isinstance((value := getattr(cls, field_name)), Case): 
            continue
        
        setattr(cls, field_name, make_dataclass(
            field_name, list(value.dict.items()), bases=(cls, )
        ))
    return cls


class Case:
    def __init__(self, **attributes):
        self.dict = attributes

    # to disable warnings
    def __call__(self, *args, **kwargs):
        pass

В результате, 17 строк кода это в принципе всё что требуется, чтобы портировать enum-ы из Rust в Python.

Результат

ИИ-модуль для следования:

@enum
class TargetChange:
    Nothing = Case()
    To = Case(target=Optional[int2])

def try_producing_target(self, subject, perception) -> TargetChange:
    if self.subject is None: 
        return TargetChange.Nothing()
    
    if abs2(sub2(subject.p, self.subject.position)) <= self.distance: 
        return TargetChange.To(None)

    if self.period.step():
        return TargetChange.To(self.subject.position)
      
    return TargetChange.Nothing()

И вот так он вызывается:

match self.follower.try_producing_target(subject, perception):
    case TargetChange.To(position): self.pather.going_to = position

Проблема решена.

P.S.

Я закинул библиотеку в PyPI и GitHub. Можете пользоваться если хотите. Также думаю о реализации Result и Option, в расте у них были полезные методы.