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