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, в расте у них были полезные методы.